Update src/loss.py

src/loss.py

@@ -96,7 +96,7 @@ class LossScheduler:
 		if timestep==A.timesteps[0]: print("resetting"); image_count+=1;count=0;
 		if A.t_prev==-1:A.xT=sample
 		A.e_prev.append(model_output)
-		if timestep==A.timesteps[-1]:
+		if timestep==51:
 			with torch.enable_grad():
 				orig_output = torch.load(f"/home/mbhat/edge-maxxing/miner/miner/latents/latent_orig_{image_count}_20.pth")
 				A.optimizer.zero_grad()