From 744682f7f9672d1a7278b5691d185e23e63fc336 Mon Sep 17 00:00:00 2001
From: gs-olive <113141689+gs-olive@users.noreply.github.com>
Date: Wed, 27 Dec 2023 20:03:24 -0800
Subject: [PATCH] fix: Switch all copies to force cast

- Fixes the case where engines return multiple copies of the same tensor
with inserted casts. This behavior is disallowed in TRT
- Added regression test case
- This fix should be removed once #2561 is merged
---
 .../dynamo/conversion/aten_ops_converters.py      |  4 ++--
 tests/py/dynamo/conversion/test_casts.py          | 15 +++++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
index fb5db527fb..8d08a0e3cd 100644
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -911,7 +911,7 @@ def aten_ops_clone_copy_dtype(
         name,
         args[0],
         kwargs.get("dtype", args[0].dtype),
-        force_layer=False,
+        force_layer=True,
     )
 
 
@@ -1043,7 +1043,7 @@ def aten_ops_sum(
             name,
             sum_,
             kwargs["output_dtype"],
-            force_layer=False,
+            force_layer=True,
         )
     else:
         return sum_
diff --git a/tests/py/dynamo/conversion/test_casts.py b/tests/py/dynamo/conversion/test_casts.py
index f17eb7b1d4..c067a0b9ad 100644
--- a/tests/py/dynamo/conversion/test_casts.py
+++ b/tests/py/dynamo/conversion/test_casts.py
@@ -75,6 +75,21 @@ def forward(self, x):
                 inputs,
             )
 
+    def test_to_copy_multiple_returns(self):
+        class ToCopyReturns(nn.Module):
+            def forward(self, x):
+                x_1 = x + 1
+                y = torch.ops.aten._to_copy.default(x_1, dtype=torch.float)
+                z = torch.ops.aten._to_copy.default(x_1, dtype=torch.float)
+                return y, z
+
+        inputs = [torch.rand((1, 3, 10))]
+        self.run_test(
+            ToCopyReturns(),
+            inputs,
+            precision=torch.float,
+        )
+
 
 if __name__ == "__main__":
     run_tests()