Merge pull request #1175 from shaltielshmid/tensor-load-cpu

Make sure to handle tensor copy when moving to cpu before loading

Author: NiklasGustafsson

RELEASENOTES.md

@@ -24,6 +24,7 @@ __Bug Fixes__:
 #1170 : Calling `torch.nn.rnn.utils.pad_packed_sequence` with a CUDA tensor and unsorted_indices threw an error
 #1172 : `optim.LoadStateDict` from an existing `StateDictionary` updated to make sure to copy value and to the right device.
 #1176 : When specific `Optimizers` load in a conditional tensor, made sure to copy to the right device.
+#1174 : Loading CUDA tensor from stream threw an error
 
 ## NuGet Version 0.101.2
 

src/TorchSharp/Tensor/TensorExtensionMethods.cs

@@ -463,9 +463,15 @@ public static void Load(this Tensor tensor, System.IO.BinaryReader reader, bool
 
             if (!skip) {
                 var device = tensor.device;
-                if (device.type != DeviceType.CPU) tensor.to(CPU);
-                tensor.bytes = bytes;
-                tensor.to(device);
+                if (device.type != DeviceType.CPU) {
+                    using var copy = tensor.to(CPU);
+                    copy.bytes = bytes;
+                    using var moved = copy.to(device);
+                    tensor.set_(moved);
+                }
+                else {
+                    tensor.bytes = bytes;
+                }
             }
         }
 
@@ -532,9 +538,10 @@ public static void Load(ref Tensor? tensor, System.IO.BinaryReader reader, bool
 
             if (!skip) {
                 var device = tensor.device;
-                if (device.type != DeviceType.CPU) tensor.to(CPU);
+                if (device.type != DeviceType.CPU)
+                    tensor = tensor.to(CPU, disposeAfter: true); 
                 tensor.bytes = bytes;
-                tensor.to(device);
+                tensor = tensor.to(device, disposeAfter: true);
             }
         }
 

test/TorchSharpTest/TestTorchTensorBugs.cs

@@ -1456,5 +1456,26 @@ public void Validate1172_Device()
                 Assert.Equal(DeviceType.CUDA, (optim2.state_dict().State[0] as Adam.State)!.exp_avg.device.type);
             }
         }
+
+        [Fact]
+        public void Validate1174()
+        {
+            if (torch.cuda.is_available()) {
+                var tensor1 = torch.ones(10, ScalarType.Int32, device: torch.CUDA);
+
+                // Save to memory stream
+                using var ms = new MemoryStream();
+                tensor1.Save(new BinaryWriter(ms, System.Text.Encoding.UTF8, true));
+                ms.Position = 0;
+
+                var tensor2 = torch.zeros(10, ScalarType.Int32, device: torch.CUDA);
+                // This used to throw an error, trying to load bytes onto a cuda tensor
+                tensor2.Load(new BinaryReader(ms));
+
+                // Make sure tensor 2 is still on CUDA with the right values
+                Assert.Equal(DeviceType.CUDA, tensor2.device.type);
+                Assert.True(Enumerable.SequenceEqual(tensor2.data<int>(), Enumerable.Repeat(1, 10)));
+            }
+        }
     }
 }