From 54d5e86981fa5425a3d1cb6cc13bfa1605523010 Mon Sep 17 00:00:00 2001 From: smrkatte <105341386+smrkatte@users.noreply.github.com> Date: Tue, 2 Aug 2022 18:32:58 -0700 Subject: [PATCH] Add cast before copy for dissimilar scalar type (#12391) * Add proper cast/copy callflow for ORT and non-ORT devices --- orttraining/orttraining/eager/ort_aten.cpp | 30 ++++++++++++++-------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/orttraining/orttraining/eager/ort_aten.cpp b/orttraining/orttraining/eager/ort_aten.cpp index b8d8b87dc4..e3c00a9884 100644 --- a/orttraining/orttraining/eager/ort_aten.cpp +++ b/orttraining/orttraining/eager/ort_aten.cpp @@ -663,20 +663,28 @@ at::Tensor& copy_( auto& invoker = GetORTInvoker(self.device().type() == at::kORT ? self.device() : src.device()); const auto ort_src = create_ort_value(invoker, src); auto ort_self = create_ort_value(invoker, self); + if (self.scalar_type() != src.scalar_type()) { - // invoke cast first - std::vector ort_cast_output(1); - onnxruntime::NodeAttributes attrs(1); - attrs["to"] = create_ort_attribute( - "to", (int64_t)GetONNXTensorProtoDataType(self.scalar_type()), at::kLong); + if(src.device().type() != at::kORT) + { + // invoke cast first and then copy for non-ORT device types + auto val = at::native::to(src, self.scalar_type()); + const auto ort_val = create_ort_value(invoker, val); + copy(invoker, ort_val, ort_self); + }else{ + // For ORT device type, the cast operation will perform the copy as well + std::vector ort_cast_output(1); + ort_cast_output[0] = ort_self; + onnxruntime::NodeAttributes attrs(1); + attrs["to"] = create_ort_attribute( + "to", (int64_t)GetONNXTensorProtoDataType(self.scalar_type()), at::kLong); - auto status = invoker.Invoke("Cast", - {std::move(ort_src)}, - ort_cast_output, &attrs); + auto status = invoker.Invoke("Cast", + {std::move(ort_src)}, + ort_cast_output, &attrs); - CHECK_STATUS(status); - - copy(invoker, ort_cast_output[0], ort_self); + CHECK_STATUS(status); + } } else { copy(invoker, ort_src, ort_self); }