fix precommit errors

Author: vahid0001

mmlearn/modules/ema.py

@@ -130,10 +130,6 @@ def restore(self, model: torch.nn.Module) -> torch.nn.Module:
         model.load_state_dict(d, strict=False)
         return model
 
-    # def state_dict(self) -> dict[str, Any]:
-    #     """Return the state dict of the model."""
-    #     return self.model.state_dict()  # type: ignore[no-any-return]
-
     @staticmethod
     def get_annealed_rate(
         start: float,

mmlearn/modules/encoders/vision.py

@@ -2,7 +2,7 @@
 
 import math
 from functools import partial
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
 
 import timm
 import torch
@@ -284,7 +284,6 @@ def __init__(
         # Weight Initialization
         self.init_std = init_std
         self.apply(self._init_weights)
-        self.fix_init_weight()
 
     def fix_init_weight(self) -> None:
         """Fix initialization of weights by rescaling them according to layer depth."""
@@ -493,7 +492,6 @@ def __init__(
         self.init_std = init_std
         trunc_normal_(self.mask_token, std=self.init_std)
         self.apply(self._init_weights)
-        # self.fix_init_weight()
 
     def fix_init_weight(self) -> None:
         """Fix initialization of weights by rescaling them according to layer depth."""
@@ -567,9 +565,12 @@ def forward(
         return self.predictor_proj(x)
 
 
-@store(
-    group="modules/encoders",
-    provider="mmlearn",
+@cast(
+    VisionTransformerPredictor,
+    store(
+        group="modules/encoders",
+        provider="mmlearn",
+    ),
 )
 def vit_predictor(**kwargs: Any) -> VisionTransformerPredictor:
     """
@@ -585,9 +586,12 @@ def vit_predictor(**kwargs: Any) -> VisionTransformerPredictor:
     )
 
 
-@store(
-    group="modules/encoders",
-    provider="mmlearn",
+@cast(
+    VisionTransformer,
+    store(
+        group="modules/encoders",
+        provider="mmlearn",
+    ),
 )
 def vit_tiny(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
     """
@@ -610,9 +614,12 @@ def vit_tiny(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
     )
 
 
-@store(
-    group="modules/encoders",
-    provider="mmlearn",
+@cast(
+    VisionTransformer,
+    store(
+        group="modules/encoders",
+        provider="mmlearn",
+    ),
 )
 def vit_small(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
     """
@@ -635,9 +642,12 @@ def vit_small(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
     )
 
 
-@store(
-    group="modules/encoders",
-    provider="mmlearn",
+@cast(
+    VisionTransformer,
+    store(
+        group="modules/encoders",
+        provider="mmlearn",
+    ),
 )
 def vit_base(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
     """
@@ -660,9 +670,12 @@ def vit_base(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
     )
 
 
-@store(
-    group="modules/encoders",
-    provider="mmlearn",
+@cast(
+    VisionTransformer,
+    store(
+        group="modules/encoders",
+        provider="mmlearn",
+    ),
 )
 def vit_large(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
     """
@@ -685,9 +698,12 @@ def vit_large(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
     )
 
 
-@store(
-    group="modules/encoders",
-    provider="mmlearn",
+@cast(
+    VisionTransformer,
+    store(
+        group="modules/encoders",
+        provider="mmlearn",
+    ),
 )
 def vit_huge(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
     """
@@ -710,9 +726,12 @@ def vit_huge(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
     )
 
 
-@store(
-    group="modules/encoders",
-    provider="mmlearn",
+@cast(
+    VisionTransformer,
+    store(
+        group="modules/encoders",
+        provider="mmlearn",
+    ),
 )
 def vit_giant(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
     """

mmlearn/tasks/ijepa_pretraining.py

@@ -154,7 +154,7 @@ def _shared_step(
 
         return loss
 
-    def configure_optimizers(self):
+    def configure_optimizers(self) -> Dict[str, Any]:
         """Configure the optimizer and learning rate scheduler."""
         weight_decay_value = 0.05  # Desired weight decay
 
@@ -194,7 +194,7 @@ def configure_optimizers(self):
             },
         ]
 
-        optimizer = torch.optim.AdamW(parameters)
+        optimizer = torch.optim.AdamW(parameters, lr=0.001)
 
         # Instantiate the learning rate scheduler if provided
         lr_scheduler = None