Implement 'to' on ScriptModules

pytorch · Dec 18, 2018 · 10fc5b3 · 10fc5b3
1 parent bb9b7de
commit 10fc5b3
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -1471,6 +1471,11 @@ def forward(self, x):
         traced_model.cpu()
         cpu_out = traced_model(x.float())
         self.assertEqual(cpu_out, cuda_out)
+        traced_model.to('cuda')
+        cuda_out = traced_model(x.float().cuda())
+        traced_model.to('cpu')
+        cpu_out = traced_model(x.float())
+        self.assertEqual(cpu_out, cuda_out)
         traced_model.double()
 
         # state_dict + load_state_dict

diff --git a/torch/jit/__init__.py b/torch/jit/__init__.py
@@ -1267,7 +1267,7 @@ def _get_methods(cls):
 
 _compiled_methods_whitelist = {
     'forward', 'register_buffer', 'register_parameter', 'add_module',
-    '_apply', 'apply', 'cuda', 'cpu', 'type', 'float', 'double', 'half',
+    '_apply', 'apply', 'cuda', 'cpu', 'to', 'type', 'float', 'double', 'half',
     'state_dict', 'load_state_dict', '_load_from_state_dict',
     '_named_members', 'parameters', 'named_parameters',
     'buffers', 'named_buffers', 'children', 'named_children', 'modules',