[dtensor][BE][1/N] fix DTensor Ops test (#113104)

XilunWu · pytorchmergebot · commit 56e514aefb1d · 2023-11-07T18:23:14.000Z
**Summary**: dtensor_ops test has helper function `assert_ref_dtensor_equal` which was written as expecting a `DTensor` argument `dtensor_rs` but actually receives a `torch.Tensor` in test. This PR removes the `to_local()` call on that object since it's actually a `torch.Tensor`. This PR is a part of internal task [T169242924](https://www.internalfb.com/intern/tasks/?t=169242924) for better engineering. **Test**: `pytest test/distributed/_tensor/test_dtensor_ops.py` `pytest test/distributed/_tensor/test_dtensor_ops.py -s -k baddbmm` Pull Request resolved: #113104 Approved by: https://github.com/wanchaol
diff --git a/test/distributed/_tensor/test_dtensor_ops.py b/test/distributed/_tensor/test_dtensor_ops.py
@@ -602,7 +602,7 @@ def assert_ref_dtensor_equal(self, dtensor_rs, rs):
                 f"dtensor requires_grad: {dtensor_r.requires_grad}",
             )
 
-            self.assertEqualOnRank(dtensor_r.to_local(), r)
+            self.assertEqualOnRank(dtensor_r, r)
 
     def run_dtensor_crossref(self, func, args, kwargs):
         to_dtensor = DTensorConverter(self.mesh, args, kwargs)
