is_xla_tensor=True in data parallel

Lightning-AI · Aug 17, 2020 · 617276e · 617276e
1 parent f833a09
commit 617276e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_lightning/trainer/distrib_data_parallel.py b/pytorch_lightning/trainer/distrib_data_parallel.py
@@ -404,7 +404,7 @@ def transfer_distrib_spawn_state_on_fit_end(self, model, mp_queue, results):
             last_path = None
             if not self.testing and best_model_path is not None and len(best_model_path) > 0:
                 last_path = re.sub('.ckpt', '.tmp_end.ckpt', best_model_path)
-                atomic_save(model.state_dict(), last_path, is_xla_tensor=self.use_tpu)
+                atomic_save(model.state_dict(), last_path, is_xla_tensor=True)
             mp_queue.put(last_path)
 
     def save_spawn_weights(self, model):