I keep running out of memory when running GaussianDiffusion1D. Is there a way to change the batch size?

[RainbowPiBubbles]

model = Unet1D( dim = 64, dim_mults = (1, 2, 4, 8), channels = 1, )

diffusion = GaussianDiffusion1D( model, seq_length = 128, timesteps = 250, objective = 'pred_v', )

training_seq = tenenc

loss = diffusion(training_seq) #Has a shape of (1200, 1, 128) loss.backward()

#after a lot of training

sampled_seq = diffusion.sample(batch_size = 1) sampled_seq.shape # (1, 1, 128)

[Adrian744]

The Trainer Class has the attribute "train_batch_size". You can simply pass the batch_size you want.