Implement 'to' on ScriptModules (#15340)

Summary:
Following #6008
Fixes "Implement 'to' on ScriptModules #7354"

cc zdevito
Pull Request resolved: https://github.com/pytorch/pytorch/pull/15340

Differential Revision: D13506646

Pulled By: zdevito

fbshipit-source-id: 318fea2e8e51a37ce9844efa4c8db67d45a66317

diff --git a/test/test_jit.py b/test/test_jit.py
index 82f89ca..84daa2e 100644 (file)
--- a/test/test_jit.py
+++ b/test/test_jit.py
@@ -1478,6 +1478,11 @@ class TestJit(JitTestCase):
         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
index ab06a34..88f52f9 100644 (file)
--- 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',