Fix batch size

train_model.py

@@ -100,16 +100,17 @@ def train_model(batch_size=4, epochs=100, save_model_interval=10, start_learning
 
     # lets train for 100 epoch for each sprite in the dataset with a random noise level
     progress_bar = tqdm(total=epochs)
+    
     for epoch in range(epochs):
         # create a noisy version of each sprite
         for batch_index in range(0, num_images, batch_size):
             progress_bar.set_description(f'epoch={epoch}, batch_index={batch_index}')
             batch_end = np.minimum(num_images, batch_index + batch_size)
             clean_images = targets[batch_index:batch_end].to(device='cuda', dtype=torch.float16)
-            clean_images = torch.reshape(clean_images, (batch_size, 12, 256, 256))
+            clean_images = torch.reshape(clean_images, ((batch_end - batch_index), 12, 256, 256))
 
             noise = torch.randn(clean_images.shape).to(device='cuda', dtype=torch.float16)
-            timesteps = torch.randint(0, model.scheduler.config.num_train_timesteps, (batch_size, ))
+            timesteps = torch.randint(0, model.scheduler.config.num_train_timesteps, (batch_end - batch_index, ))
             timesteps = timesteps.to(dtype=torch.int, device='cuda')
             noisy_images = model.scheduler.add_noise(clean_images, noise, timesteps).to(device='cuda', dtype=torch.float16)
             noise_pred = model.unet(noisy_images, timesteps, class_labels[batch_index:batch_end].to(device='cuda',dtype=torch.float16), return_dict=False)[0]
@@ -128,4 +129,4 @@ def train_model(batch_size=4, epochs=100, save_model_interval=10, start_learning
 
 
 if __name__ == '__main__':
-    train_model()
+    train_model(8)