Copy input tensors before async transfer (#5830)

* Copy input tensors before transfer * clone tensors before passing them to test case * formatting * Update test_utils.py
pytorch · Apr 22, 2024 · d39e51b · d39e51b
1 parent 47b1950
commit d39e51b
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/torch_xla/csrc/runtime/tensor_source.h b/torch_xla/csrc/runtime/tensor_source.h
@@ -53,10 +53,9 @@ class AtenSource : public TensorSource {
     at::ScalarType target_torch_type = TorchTypeFromXlaType(primitive_type());
     if (target_torch_type != tensor.type().scalarType()) {
       TORCH_LAZY_COUNTER("AtenSourceDowncasts", 1);
-      tensor_ = std::move(tensor.to(target_torch_type).contiguous());
-    } else {
-      tensor_ = std::move(tensor.contiguous());
     }
+    tensor_ = std::move(tensor.to(target_torch_type, /*non_blocking=*/false,
+                                  /*copy=*/true, at::MemoryFormat::Contiguous));
   }
 
   const void* data() const override { return tensor_.const_data_ptr(); }