check result

srossross · srossross · commit 0e77eb41ad5b · 2023-01-11T14:43:31.000-06:00
diff --git a/test/functorch/test_vmap.py b/test/functorch/test_vmap.py
@@ -350,11 +350,14 @@ def foo(x):
     def test_out_dims_normal_tensor(self):
 
         def foo(x):
-            return torch.zeros(3)
+            return torch.arange(3)
 
         tensor = torch.randn(2, 3)
-        vmap(foo)(tensor)
-        vmap(foo, out_dims=None)(tensor)
+        result = vmap(foo)(tensor)
+        self.assertEqual(result.shape, [2, 3])
+
+        result = vmap(foo, out_dims=None)(tensor)
+        self.assertEqual(result, torch.arange(3))
 
 
     def test_pytree_returns(self):
