update comment for navit 3d

lucidrains · Nov 8, 2024 · 6693d47 · 6693d47
1 parent 141239c
commit 6693d47
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/vit_pytorch/na_vit_nested_tensor.py b/vit_pytorch/na_vit_nested_tensor.py
@@ -323,3 +323,5 @@ def forward(
     ]
 
     assert v(images).shape == (5, 1000)
+
+    v(images).sum().backward()
diff --git a/vit_pytorch/na_vit_nested_tensor_3d.py b/vit_pytorch/na_vit_nested_tensor_3d.py
@@ -336,7 +336,7 @@ def forward(
 
 if __name__ == '__main__':
 
-    # works for torch 2.4
+    # works for torch 2.5
 
     v = NaViT(
         image_size = 256,
@@ -362,3 +362,5 @@ def forward(
     ]
 
     assert v(volumes).shape == (5, 1000)
+
+    v(volumes).sum().backward()
Original file line number	Diff line number	Diff line change
Expand Up		@@ -323,3 +323,5 @@ def forward(
		]

		assert v(images).shape == (5, 1000)

		v(images).sum().backward()