Distributed Training Gets Stuck?

[pjspol]

The following is part of a script that I'm trying to run using multiple GPUs (through runpod.io). Unfortunately, the training gets stuck at loss = trainer(.... I would also like to track the loss as it progresses if at all possible. I have a bit of experience with PyTorch Lightning but very little with Accelerate. Any help would be much appreciated!

from imagen_pytorch import Imagen, BaseUnet64, SRUnet256, SRUnet1024, ImagenTrainer

unet1 = BaseUnet64()

unet2 = SRUnet256()

unet3 = SRUnet1024()

imagen = Imagen(
    unets = (unet1, unet2, unet3),
    image_sizes = (64, 256, 1024),
    timesteps = 1000,
    cond_drop_prob = 0.1,
    channels = 1,
    text_embed_dim = 4096
)

trainer = ImagenTrainer(
    imagen = imagen,
    split_valid_from_train = False
).cuda()

trainer.save('/workspace/save_square1.pt')

n_epochs = 100

for i in [1, 2, 3]:
  trainer = ImagenTrainer(
    imagen = imagen,
    split_valid_from_train = False
  ).cuda()
  trainer.load('/workspace/save_square1.pt')
  for j in range(n_epochs):
    loss = trainer(images,
                  text_embeds = text_embeds,
                  unet_number = i,
                  max_batch_size = 1)
    trainer.update(unet_number = i)
  trainer.accelerator.wait_for_everyone()
  if trainer.is_main:
    trainer.save('/workspace/save_square1.pt')

[FriedRonaldo]

issue #309 might be the same problem.

TL;DR: remove "if trainer.is_main:" -- trainer.save automatically checks if it's main or not.