Update src/loss.py

src/loss.py

@@ -22,7 +22,7 @@ class LossSchedulerModel(torch.nn.Module):
         """
         batch_size, num_channels, height, width = xT.shape
         # num_channels = 1
-        num_prev_steps = e_prev.shape[0]
+        num_prev_steps = len(e_prev)
         
         # Ensure inputs match expected dimensions
         assert timestep - num_prev_steps + 1 == 0, "Mismatch between timestep and e_prev length"