update comment for navit 3d

Author: lucidrains

vit_pytorch/na_vit_nested_tensor.py

@@ -323,3 +323,5 @@ def forward(
     ]
 
     assert v(images).shape == (5, 1000)
+
+    v(images).sum().backward()

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()