Code stuck after running 1 epoch on TPU

[sumanthd17]

❓ 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

[awaelchli]

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

[sumanthd17]

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> .