question: Code stuck after running 1 epoch on TPU
context: ❓ Questions and Help
  What is your question?
  I'm trying to run the LitAutoEncoder on TPUs, but the code runs for 1 epoch and gets stuck there.
  Code
  class LitAutoEncoder(pl.LightningModule):
  
      def __init__(self, hparams):
          super().__init__()
          self.hparams = hparams
          self.encoder = nn.Sequential(
              nn.Linear(28*28, 64),
              nn.ReLU(),
              nn.Linear(64, 3)
          )
          self.decoder = nn.Sequential(
              nn.Linear(3, 64),
              nn.ReLU(),
              nn.Linear(64, 28*28)
          )
  
      def forward(self, x):
          # in lightning, forward defines the prediction/inference actions
          embedding = self.encoder(x)
          return embedding
  
      def training_step(self, batch, batch_idx):
          # training_step defined the train loop.
          # It is independent of forward
          x, y = batch
          x = x.view(x.size(0), -1)
          z = self.encoder(x)
          x_hat = self.decoder(z)
          loss = F.mse_loss(x_hat, x)
          # Logging to TensorBoard by default
          self.log('train_loss', loss)
          return loss
  
      def configure_optimizers(self):
          optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
          return optimizer
  
  dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
  train_loader = DataLoader(dataset, drop_last=True, batch_size=32)
  
  args_dict = dict(
      num_train_epochs=1,
      seed=42,
  )
  
  args = argparse.Namespace(**args_dict)
  autoencoder = LitAutoEncoder(args)
  
  train_params = dict(
      tpu_cores=8,
      progress_bar_refresh_rate=30,
  )
  
  trainer = pl.Trainer(**train_params)
  trainer.fit(autoencoder, train_loader)
  
  
  Reproducible Colab Notebook
  Notebook
  What's your environment?
  
  Colab
  Packaging pip
  pytorch-1.7
  pytorch-lightning-1.1.5
answer: Thanks Adrian,
  
  Sorry I've been editing the notebook. It's running now, but not sure what
  the issue was. My bad I ended up editing the notebook and error was not
  reproducible anymore. Will update it back to the old version.
  …
  On Sun, Jan 24, 2021 at 12:00 AM Adrian Wälchli ***@***.***> wrote:
   Could you check again? I just ran your colab and it finished both epochs.
   Maybe it's random? So far did not see random behaviour
  
   —
   You are receiving this because you authored the thread.
   Reply to this email directly, view it on GitHub
   <#5625 (comment)>,
   or unsubscribe
   <https://github.com/notifications/unsubscribe-auth/AGX3GHX6NN62CNNLKPFB6O3S3MIVXANCNFSM4WPSSWOA>
   .