from review

srossross · srossross · commit a570975d8ae4 · 2023-01-17T10:44:55.000-06:00
diff --git a/test/functorch/test_vmap.py b/test/functorch/test_vmap.py
@@ -324,7 +324,7 @@ def foo(x):
         tensor = torch.randn(2, 3)
         result = vmap(foo, out_dims=(0, None))(tensor)
         self.assertEqual(result[1], 'hello world')
-        self.assertIsInstance(result[0], torch.Tensor)
+        self.assertEqual(result[0], tensor)
 
         def foo(x):
             x.add_(1)
diff --git a/torch/_functorch/vmap.py b/torch/_functorch/vmap.py
@@ -145,7 +145,7 @@ def _maybe_remove_batch_dim(name, batched_output, vmap_level, batch_size, out_di
     # out_dim is non None
     if not isinstance(batched_output, torch.Tensor):
         raise ValueError(f'vmap({name}, ...): `{name}` must only return '
-                         f'Tensors, got type {type(batched_output)} as a return. '
+                         f'Tensors, got type {type(batched_output)}. '
                          'Did you mean to set out_dim= to None for output?')
 
     return _remove_batch_dim(batched_output, vmap_level, batch_size, out_dim)
