diff --git a/openselfsup/models/byol.py b/openselfsup/models/byol.py
index 823253f6..985971a6 100644
--- a/openselfsup/models/byol.py
+++ b/openselfsup/models/byol.py
@@ -15,9 +15,10 @@ class BYOL(nn.Module):
     Self-Supervised Learning (https://arxiv.org/abs/2006.07733)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
+        backbone (dict): Config dict for module of backbone ConvNet.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
         base_momentum (float): The base momentum coefficient for the target network.
             Default: 0.996.
diff --git a/openselfsup/models/classification.py b/openselfsup/models/classification.py
index 721f40b0..f19b6038 100644
--- a/openselfsup/models/classification.py
+++ b/openselfsup/models/classification.py
@@ -14,9 +14,9 @@ class Classification(nn.Module):
     """Simple image classification.
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
+        backbone (dict): Config dict for module of backbone ConvNet.
         with_sobel (bool): Whether to apply a Sobel filter on images. Default: False.
-        head (nn.Module): Module of loss functions.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
     """
 
diff --git a/openselfsup/models/deepcluster.py b/openselfsup/models/deepcluster.py
index b5f5cd8a..e68a73b5 100644
--- a/openselfsup/models/deepcluster.py
+++ b/openselfsup/models/deepcluster.py
@@ -18,10 +18,11 @@ class DeepCluster(nn.Module):
     of Visual Features (https://arxiv.org/abs/1807.05520)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
+        backbone (dict): Config dict for module of backbone ConvNet.
         with_sobel (bool): Whether to apply a Sobel filter on images. Default: False.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
     """
 
diff --git a/openselfsup/models/moco.py b/openselfsup/models/moco.py
index b3e8fc87..6ce83d6f 100644
--- a/openselfsup/models/moco.py
+++ b/openselfsup/models/moco.py
@@ -17,9 +17,10 @@ class MOCO(nn.Module):
     "https://github.com/facebookresearch/moco/blob/master/moco/builder.py".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
+        backbone (dict): Config dict for module of backbone ConvNet.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
         queue_len (int): Number of negative keys maintained in the queue.
             Default: 65536.
diff --git a/openselfsup/models/npid.py b/openselfsup/models/npid.py
index a61ed92b..0203919d 100644
--- a/openselfsup/models/npid.py
+++ b/openselfsup/models/npid.py
@@ -15,10 +15,11 @@ class NPID(nn.Module):
     Instance Discrimination (https://arxiv.org/abs/1805.01978)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
-        memory_bank (nn.Module): Module of memory banks.
+        backbone (dict): Config dict for module of backbone ConvNet.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
+        memory_bank (dict): Config dict for module of memory banks. Default: None.
         neg_num (int): Number of negative samples for each image. Default: 65536.
         ensure_neg (bool): If False, there is a small probability
             that negative samples contain positive ones. Default: False.
diff --git a/openselfsup/models/odc.py b/openselfsup/models/odc.py
index 9acc6332..34c03c53 100644
--- a/openselfsup/models/odc.py
+++ b/openselfsup/models/odc.py
@@ -17,11 +17,12 @@ class ODC(nn.Module):
     (https://arxiv.org/abs/2006.10645)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
+        backbone (dict): Config dict for module of backbone ConvNet.
         with_sobel (bool): Whether to apply a Sobel filter on images. Default: False.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
-        memory_bank (nn.Module): Module of memory banks.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
+        memory_bank (dict): Module of memory banks. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
     """
 
diff --git a/openselfsup/models/relative_loc.py b/openselfsup/models/relative_loc.py
index 56d23fc6..173eb7c2 100644
--- a/openselfsup/models/relative_loc.py
+++ b/openselfsup/models/relative_loc.py
@@ -15,9 +15,10 @@ class RelativeLoc(nn.Module):
     by Context Prediction (https://arxiv.org/abs/1505.05192)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
+        backbone (dict): Config dict for module of backbone ConvNet.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
     """
 
diff --git a/openselfsup/models/rotation_pred.py b/openselfsup/models/rotation_pred.py
index e1cd844f..59049aa8 100644
--- a/openselfsup/models/rotation_pred.py
+++ b/openselfsup/models/rotation_pred.py
@@ -15,8 +15,8 @@ class RotationPred(nn.Module):
     by Predicting Image Rotations (https://arxiv.org/abs/1803.07728)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
-        head (nn.Module): Module of loss functions.
+        backbone (dict): Config dict for module of backbone ConvNet.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
     """
 
diff --git a/openselfsup/models/simclr.py b/openselfsup/models/simclr.py
index c7b16587..42073596 100644
--- a/openselfsup/models/simclr.py
+++ b/openselfsup/models/simclr.py
@@ -16,9 +16,10 @@ class SimCLR(nn.Module):
     of Visual Representations (https://arxiv.org/abs/2002.05709)".
 
     Args:
-        backbone (nn.Module): Module of backbone ConvNet.
-        neck (nn.Module): Module of deep features to compact feature vectors.
-        head (nn.Module): Module of loss functions.
+        backbone (dict): Config dict for module of backbone ConvNet.
+        neck (dict): Config dict for module of deep features to compact feature vectors.
+            Default: None.
+        head (dict): Config dict for module of loss functions. Default: None.
         pretrained (str, optional): Path to pre-trained weights. Default: None.
     """