required_changes

lightly-ai · Oct 24, 2024 · fb70bc8 · fb70bc8
1 parent 328daad
commit fb70bc8
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/lightly/models/modules/ijepa.py b/lightly/models/modules/ijepa.py
@@ -117,14 +117,14 @@ def forward(self, x, masks_x, masks):
 
         Args:
             x:
-                Input tensor.
+                Input tensor of shape (batch_size, sequence_length, feature_dim).
             masks_x:
-                Mask indices for the input tensor.
+                Mask indices for the input tensor of shape (batch_size, num_patches).
             masks:
-                Mask indices for the predicted tokens.
+                Mask indices for the predicted tokens of shape (batch_size, num_patches).
 
         Returns:
-            The predicted output tensor.
+            The predicted output tensor of shape (batch_size, num_patches, output_dim).
         """
         assert (masks is not None) and (
             masks_x is not None

diff --git a/lightly/models/modules/ijepa_timm.py b/lightly/models/modules/ijepa_timm.py
@@ -103,14 +103,14 @@ def forward(
 
         Args:
             x:
-                Input tensor.
+                Input tensor of shape (batch_size, sequence_length, feature_dim).
             masks_x:
-                Mask indices for the input tensor.
+                Mask indices for the input tensor of shape (batch_size, num_patches).
             masks:
-                Mask indices for the predicted tokens.
+                Mask indices for the predicted tokens of shape (batch_size, num_patches).
 
         Returns:
-            The predicted output tensor.
+            The predicted output tensor of shape (batch_size, num_patches, output_dim).
         """
 
         assert (masks is not None) and (