Fixes #902

pytorch_lightning/core/lightning.py

@@ -16,6 +16,13 @@
 from pytorch_lightning.core.memory import ModelSummary
 from pytorch_lightning.overrides.data_parallel import LightningDistributedDataParallel
 
+try:
+ import torch_xla.core.xla_model as xm
+ XLA_AVAILABLE = True
+
+except ImportError:
+ XLA_AVAILABLE = False
+
 
 class LightningModule(ABC, GradInformation, ModelIO, ModelHooks):
 
@@ -798,7 +805,9 @@ def optimizer_step(self, current_epoch, batch_idx, optimizer, optimizer_idx, sec
  optimizer.zero_grad()
 
  """
- if isinstance(optimizer, torch.optim.LBFGS):
+ if self.trainer.use_tpu and XLA_AVAILABLE:
+ xm.optimizer_step(optimizer)
+ elif isinstance(optimizer, torch.optim.LBFGS):
  optimizer.step(second_order_closure)
  else:
  optimizer.step()

pytorch_lightning/trainer/training_loop.py

@@ -168,15 +168,9 @@ def training_step(self, batch, batch_idx):
 except ImportError:
  APEX_AVAILABLE = False
 
-try:
- import torch_xla.core.xla_model as xm
-
- XLA_AVAILABLE = True
-except ImportError:
- XLA_AVAILABLE = False
-
 try:
  import torch_xla.distributed.parallel_loader as xla_pl
+ import torch_xla.core.xla_model as xm
 
  XLA_AVAILABLE = True
 
@@ -600,11 +594,8 @@ def optimizer_closure():
  # override function to modify this behavior
  model = self.get_model()
  with self.profiler.profile('optimizer_step'):
- if self.use_tpu:
- xm.optimizer_step(optimizer)
- else:
- model.optimizer_step(self.current_epoch, batch_idx,
- optimizer, opt_idx, optimizer_closure)
+ model.optimizer_step(self.current_epoch, batch_idx,
+ optimizer, opt_idx, optimizer_closure)
 
  # calculate running loss for display
  self.running_loss.append(self.batch_loss_value)