Fixed logic to keep track of which device ParameterDict and ParameterList modules are held on.

Author: Niklas Gustafsson

src/TorchSharp/JIT/ScriptModule.cs

@@ -190,9 +190,6 @@ protected internal override nn.Module _to(DeviceType deviceType, int deviceIndex
                     return this;
                 }
 
-                private DeviceType _deviceType = DeviceType.CPU;
-                private int _deviceIndex = -1;
-
                 /// <summary>
                 /// Convert the parameters and buffers.
                 /// </summary>

src/TorchSharp/NN/Module.cs

@@ -203,8 +203,8 @@ protected internal virtual Module _to(DeviceType deviceType, int deviceIndex, bo
                     return this;
                 }
 
-                private DeviceType _deviceType = DeviceType.CPU;
-                private int _deviceIndex = -1;
+                protected DeviceType _deviceType = DeviceType.CPU;
+                protected int _deviceIndex = -1;
 
                 /// <summary>
                 /// Convert the parameters and buffers.

src/TorchSharp/NN/ParameterDict.cs

@@ -98,6 +98,11 @@ protected override void _toEpilog(torch.ScalarType? dtype, torch.Device device,
                     _list[i] = (name, p);
                     _dict[name] = p;
                 }
+
+                if (device is not null) {
+                    _deviceType = device.type;
+                    _deviceIndex = device.index;
+                }
             }
 
             /// <summary>

src/TorchSharp/NN/ParameterList.cs

@@ -90,6 +90,11 @@ protected override void _toEpilog(torch.ScalarType? dtype, torch.Device? device,
 
                     _list[i] = p;
                 }
+
+                if (device is not null) {
+                    _deviceType = device.type;
+                    _deviceIndex = device.index;
+                }
             }
 
             private bool _registered = false;

test/TorchSharpTest/NN.cs

@@ -3170,6 +3170,32 @@ public void TestCustomModuleWithDeviceMove()
                 // Reset and then try again with moving back to CPU
                 module.zero_grad();
 
+                // Try moving back to CPU
+                module.to(torch.CPU);
+                x = torch.randn(2, 2);
+                y = torch.randn(2);
+                torch.nn.functional.mse_loss(module.call(x), y).backward();
+                foreach (var (pName, parm) in module.named_parameters()) {
+                    var grad = parm.grad;
+                    Assert.NotNull(grad);
+                }
+            }
+            if (torch.mps_is_available()) {
+                var module = new TestModule1(torch.randn(2, 2), true);
+
+                // Move the device to MPS, and make sure gradients are calculated for all the parameters
+                module.to(torch.MPS);
+                var x = torch.randn(2, 2, device: torch.MPS);
+                var y = torch.randn(2, device: torch.MPS);
+                torch.nn.functional.mse_loss(module.call(x), y).backward();
+                foreach (var (pName, parm) in module.named_parameters()) {
+                    var grad = parm.grad;
+                    Assert.NotNull(grad);
+                }
+
+                // Reset and then try again with moving back to CPU
+                module.zero_grad();
+
                 // Try moving back to CPU
                 module.to(torch.CPU);
                 x = torch.randn(2, 2);