diff --git a/.github/workflows/latest-changes.yml b/.github/workflows/latest-changes.yml
index 52211ab2..c124d030 100644
--- a/.github/workflows/latest-changes.yml
+++ b/.github/workflows/latest-changes.yml
@@ -24,6 +24,5 @@ jobs:
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
           latest_changes_file: docs/CHANGELOG.md
-          template_file: ./.github/workflows/release-notes.jinja2
-          latest_changes_header: '## 0.0.1\n'
+          latest_changes_header: '## 0.0.2\n\n'
           debug_logs: true
diff --git a/.github/workflows/release-notes.jinja2 b/.github/workflows/release-notes.jinja2
deleted file mode 100644
index 0d02bd69..00000000
--- a/.github/workflows/release-notes.jinja2
+++ /dev/null
@@ -1 +0,0 @@
-* This changed: {{pr.title}}. Done by [ {{pr.user.login}}]({{pr.user.html_url}}). Check the [Pull Request {{pr.number}} with the changes and stuff]({{pr.html_url}}). now back to code. 🤓\n
diff --git a/Makefile b/Makefile
index 59aff998..35f8b83c 100644
--- a/Makefile
+++ b/Makefile
@@ -1,8 +1,8 @@
 build-docs:
 	cp README.md docs/index.md
 
-docs-serve:
-	mkdocs serve
+docsserve:
+	mkdocs serve --dirtyreload
 
 test:
 	python tests/__init__.py
diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index 95fa6600..cc062d05 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -1,12 +1,14 @@
 # Release Notes
 
 ## 0.0.1
-* This changed: 📝 update example and documentation. Done by [ aniketmaurya](https://github.com/aniketmaurya). Check the [Pull Request 20 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/20). now back to code. 🤓\n
-* This changed: :tada::sparkles: First Release - v0.0.1 - Refactor API & tested Python 3.7+. Done by [ aniketmaurya](https://github.com/aniketmaurya). Check the [Pull Request 18 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/18). now back to code. 🤓
-* This changed: Adding example notebook for AutoSummarization. Done by [the GitHub user gagan3012](https://github.com/gagan3012). Check the [Pull Request 19 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/19). now back to code. 🤓
-* This changed: Adding text summarisation. Done by [the GitHub user gagan3012](https://github.com/gagan3012). Check the [Pull Request 14 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/14). now back to code. 🤓
-* This changed: add codecov CI. Done by [the GitHub user aniketmaurya](https://github.com/aniketmaurya). Check the [Pull Request 15 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/15). now back to code. 🤓
-* This changed: 📚 update documentation - added citation, acknowledgments, docstrings automation. Done by [the GitHub user aniketmaurya](https://github.com/aniketmaurya). Check the [Pull Request 13 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/13). now back to code. 🤓
+* 📝 update example and documentation. Done by [ aniketmaurya](https://github.com/aniketmaurya). Check the [Pull Request 20 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/20).
+* :tada::sparkles: First Release - v0.0.1 - Refactor API & tested Python 3.7+. Done by [ aniketmaurya](https://github.com/aniketmaurya). Check the [Pull Request 18 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/18).
+* Adding example notebook for AutoSummarization. Done by [the GitHub user gagan3012](https://github.com/gagan3012). Check the [Pull Request 19 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/19).
+* Adding text summarisation. Done by [the GitHub user gagan3012](https://github.com/gagan3012). Check the [Pull Request 14 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/14).
+* add codecov CI. Done by [the GitHub user aniketmaurya](https://github.com/aniketmaurya). Check the [Pull Request 15 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/15).
+* 📚 update documentation - added citation, acknowledgments, docstrings automation. Done by [the GitHub user aniketmaurya](https://github.com/aniketmaurya). Check the [Pull Request 13 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/13).
 * Refactor API Design, CI & Docs PR [#10](https://github.com/gradsflow/gradsflow/pull/10) by [@aniketmaurya](https://github.com/aniketmaurya).
 * auto docstring. PR [#7](https://github.com/gradsflow/gradsflow/pull/7) by [@aniketmaurya](https://github.com/aniketmaurya).
 * Add AutoImageClassifier. PR [#1](https://github.com/gradsflow/gradsflow/pull/1) by [@aniketmaurya](https://github.com/aniketmaurya).
+
+## 0.0.2
diff --git a/gradsflow/core/autoclassifier.py b/gradsflow/core/autoclassifier.py
index 2c0f1e1c..ce88ed6b 100644
--- a/gradsflow/core/autoclassifier.py
+++ b/gradsflow/core/autoclassifier.py
@@ -1,3 +1,17 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
 from abc import abstractmethod
 from typing import Dict, List, Optional, Union
 
@@ -12,7 +26,7 @@
 
 
 class AutoClassifier(AutoModel):
-    """Base Class for Auto Classification Hyperparameter search"""
+    """Implements `AutoModel` for classification tasks."""
 
     DEFAULT_BACKBONES = []
 
diff --git a/gradsflow/core/automodel.py b/gradsflow/core/automodel.py
index bea62dcd..922d99b5 100644
--- a/gradsflow/core/automodel.py
+++ b/gradsflow/core/automodel.py
@@ -1,3 +1,17 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
 from abc import abstractmethod
 from typing import Dict, Optional, Union
 
@@ -5,20 +19,36 @@
 import pytorch_lightning as pl
 import torch
 from flash import DataModule
+from loguru import logger
 from optuna.integration import PyTorchLightningPruningCallback
 
-from gradsflow.logging import logger
 from gradsflow.utility.common import module_to_cls_index
+from gradsflow.utility.optuna import is_best_trial
 
 
 class AutoModel:
     """
     Creates Optuna instance, defines methods required for hparam search
+
+    Args:
+        datamodule [flash.DataModule]: DataModule from Flash or PyTorch Lightning
+        max_epochs [int]: Maximum number of epochs for which model will train
+        optimization_metric [str]: Value on which hyperparameter search will run.
+        By default, it is `val_accuracy`.
+        n_trials [int]: Number of trials for HPO
+        suggested_conf [Dict]: Any extra suggested configuration
+        timeout [int]: HPO will stop after timeout
+        prune [bool]: Whether to stop unpromising training.
+        optuna_confs [Dict]: Optuna configs
+        best_trial [bool]: If true model will be loaded with best weights from HPO otherwise
+        a best trial model without trained weights will be created.
     """
 
     OPTIMIZER_INDEX = module_to_cls_index(torch.optim, True)
     DEFAULT_OPTIMIZERS = ["adam", "sgd"]
     DEFAULT_LR = (1e-5, 1e-1)
+    _BEST_MODEL = "best_model"
+    _CURRENT_MODEL = "current_model"
 
     def __init__(
         self,
@@ -30,6 +60,7 @@ def __init__(
         timeout: int = 600,
         prune: bool = True,
         optuna_confs: Optional[Dict] = None,
+        best_trial: bool = True,
     ):
 
         self._pruner: optuna.pruners.BasePruner = (
@@ -37,6 +68,7 @@ def __init__(
         )
         self.datamodule = datamodule
         self.n_trials = n_trials
+        self.best_trial = best_trial
         self.model: Union[torch.nn.Module, pl.LightningModule, None] = None
         self.max_epochs = max_epochs
         self.timeout = timeout
@@ -96,6 +128,7 @@ def _objective(
         )
         trial_confs = self._get_trial_hparams(trial)
         model = self.build_model(**trial_confs)
+        trial.set_user_attr(key="current_model", value=model)
         hparams = dict(model=model.hparams)
         trainer.logger.log_hyperparams(hparams)
         trainer.fit(model, datamodule=self.datamodule)
@@ -103,11 +136,27 @@ def _objective(
         logger.debug(trainer.callback_metrics)
         return trainer.callback_metrics[self.optimization_metric].item()
 
+    def callback_best_trial(self, study: optuna.Study, trial: optuna.Trial) -> None:
+        if is_best_trial(study, trial):
+            study.set_user_attr(
+                key=self._BEST_MODEL, value=trial.user_attrs[self._CURRENT_MODEL]
+            )
+
     def hp_tune(self):
         """
         Search Hyperparameter and builds model with the best params
         """
+        callbacks = []
+        if self.best_trial:
+            callbacks.append(self.callback_best_trial)
         self._study.optimize(
-            self._objective, n_trials=self.n_trials, timeout=self.timeout
+            self._objective,
+            n_trials=self.n_trials,
+            timeout=self.timeout,
+            callbacks=callbacks,
         )
-        self.model = self.build_model(**self._study.best_params)
+
+        if self.best_trial:
+            self.model = self._study.user_attrs[self._BEST_MODEL]
+        else:
+            self.model = self.build_model(**self._study.best_params)
diff --git a/gradsflow/logging.py b/gradsflow/logging.py
deleted file mode 100644
index fcc3e275..00000000
--- a/gradsflow/logging.py
+++ /dev/null
@@ -1 +0,0 @@
-from loguru import logger
diff --git a/gradsflow/utility/optuna.py b/gradsflow/utility/optuna.py
new file mode 100644
index 00000000..89e49566
--- /dev/null
+++ b/gradsflow/utility/optuna.py
@@ -0,0 +1,33 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import optuna
+
+
+def is_best_trial(study: optuna.Study, trial: optuna.Trial) -> bool:
+    if study.best_trial.number == trial.number:
+        return True
+    return False
diff --git a/pyproject.toml b/pyproject.toml
index 9773492a..21d97ddc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,11 +27,11 @@ classifiers = [
 ]
 requires-python = ">=3.7"
 requires = [
-    "optuna==2.9.1",
-    "smart_open==5.1.0",
+    "optuna==2.9",
+    "smart_open==5.1",
     "lightning-flash[all]==0.4.0",
     "pytorch-lightning==1.4.0",
-    "loguru==0.5"
+    "loguru~=0.5"
 ]
 
 [tool.flit.metadata.urls]
diff --git a/tests/__init__.py b/tests/__init__.py
index e69de29b..f775d5ed 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -0,0 +1,13 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
diff --git a/tests/__main__.py b/tests/__main__.py
index d2b9b1b9..9fd43a7f 100644
--- a/tests/__main__.py
+++ b/tests/__main__.py
@@ -1,3 +1,17 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
 from flash.core.data.utils import download_data
 
 download_data("https://pl-flash-data.s3.amazonaws.com/hymenoptera_data.zip", "./data")
diff --git a/tests/core/__init__.py b/tests/core/__init__.py
index e69de29b..f775d5ed 100644
--- a/tests/core/__init__.py
+++ b/tests/core/__init__.py
@@ -0,0 +1,13 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
diff --git a/tests/core/test_autoclassifier.py b/tests/core/test_autoclassifier.py
index e69de29b..f775d5ed 100644
--- a/tests/core/test_autoclassifier.py
+++ b/tests/core/test_autoclassifier.py
@@ -0,0 +1,13 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
diff --git a/tests/core/test_automodel.py b/tests/core/test_automodel.py
index e8ede67a..909721ca 100644
--- a/tests/core/test_automodel.py
+++ b/tests/core/test_automodel.py
@@ -1,3 +1,17 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
 import pytest
 from flash.image import ImageClassificationData
 
diff --git a/tests/tasks/__init__.py b/tests/tasks/__init__.py
index e69de29b..f775d5ed 100644
--- a/tests/tasks/__init__.py
+++ b/tests/tasks/__init__.py
@@ -0,0 +1,13 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
diff --git a/tests/tasks/test_image.py b/tests/tasks/test_image.py
index f51d41ed..85fc6696 100644
--- a/tests/tasks/test_image.py
+++ b/tests/tasks/test_image.py
@@ -1,3 +1,17 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
 import pytest
 import torch
 from flash.image import ImageClassificationData, ImageClassifier
diff --git a/tests/tasks/test_summarization.py b/tests/tasks/test_summarization.py
index ebcd3b86..20097b2c 100644
--- a/tests/tasks/test_summarization.py
+++ b/tests/tasks/test_summarization.py
@@ -1,7 +1,21 @@
-from gradsflow.autotasks import AutoSummarization
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
 
 from unittest.mock import MagicMock
 
+from gradsflow.autotasks import AutoSummarization
+
 
 def test_build_model():
     datamodule = MagicMock()
diff --git a/tests/tasks/test_text.py b/tests/tasks/test_text.py
index 148ad7ec..0d8485ae 100644
--- a/tests/tasks/test_text.py
+++ b/tests/tasks/test_text.py
@@ -1,3 +1,17 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
 from unittest.mock import MagicMock
 
 from gradsflow.autotasks import AutoTextClassifier
diff --git a/tests/utility/__init__.py b/tests/utility/__init__.py
new file mode 100644
index 00000000..f775d5ed
--- /dev/null
+++ b/tests/utility/__init__.py
@@ -0,0 +1,13 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
diff --git a/tests/utility/optuna/__init__.py b/tests/utility/optuna/__init__.py
new file mode 100644
index 00000000..f775d5ed
--- /dev/null
+++ b/tests/utility/optuna/__init__.py
@@ -0,0 +1,13 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
diff --git a/tests/utility/optuna/test_optuna.py b/tests/utility/optuna/test_optuna.py
new file mode 100644
index 00000000..fff33d7a
--- /dev/null
+++ b/tests/utility/optuna/test_optuna.py
@@ -0,0 +1,42 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from unittest.mock import MagicMock
+
+from gradsflow.utility.optuna import is_best_trial
+
+
+def test_is_best_trial():
+    trial = MagicMock()
+
+    study = MagicMock()
+    study.best_trial.number = 1
+    trial.number = 1
+
+    assert is_best_trial(study, trial)
+
+    trial.number = 2
+    assert is_best_trial(study, trial) is False
diff --git a/tests/utility/test_common.py b/tests/utility/test_common.py
index fb47764a..e06da7b8 100644
--- a/tests/utility/test_common.py
+++ b/tests/utility/test_common.py
@@ -1,3 +1,17 @@
+#  Copyright (c) 2021 GradsFlow. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
 import torch
 
 from gradsflow.utility.common import (