Update src/loss.py

src/loss.py

@@ -110,7 +110,8 @@ class SchedulerWrapper:
 		D=sample;E=model_output;A=timestep; 
 		if timestep==0: print("resetting"); image_count+=1;count=0
 		if True:
-			C=B.scheduler.step(E,A,D,**F);A=A.tolist(); if timestep+1==len(B.timesteps): torch.save(model_output, f"/home/mbhat/edge-maxxing/miner/miner/latents/latent_{count}.pth")
+			C=B.scheduler.step(E,A,D,**F);A=A.tolist(); 
+			if timestep+1==len(B.timesteps): torch.save(model_output, f"/home/mbhat/edge-maxxing/miner/miner/latents/latent_{count}.pth")
 			if A not in B.catch_x:B.catch_x[A]=[];B.catch_e[A]=[];B.catch_x_[A]=[]
 			B.catch_x[A].append(D.clone().detach().cpu());B.catch_e[A].append(E.clone().detach().cpu());B.catch_x_[A].append(C[0].clone().detach().cpu());return C
 		else:C=B.loss_scheduler.step(E,A,D,**F);return C