Convert entire model to jax xla (#8330)

pytorch · Oct 29, 2024 · 2592a68 · 2592a68
1 parent 3d1d1b9
commit 2592a68
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/experimental/torch_xla2/torch_xla2/tensor.py b/experimental/torch_xla2/torch_xla2/tensor.py
@@ -452,12 +452,8 @@ def __exit__(self, *exc):
 
     def _move_one_value(self, val):
       if isinstance(val, torch.nn.Module):
-        state_dict = self.to_xla(val.state_dict())
-        val.load_state_dict(state_dict, assign=True)
-        # Non-persistent buffers are not in state_dict
-        for b_name, buffer in val.named_buffers():
-          setattr(val, b_name, self.to_xla(buffer))
-        return val
+        with self:
+          return val.to('jax')
       if isinstance(val, XLATensor2):
         return val
       if isinstance(val, torch.Tensor):