Update src/loss.py

src/loss.py

@@ -1,6 +1,6 @@
+_A=None
 count=0
 image_count=0
-_A=None
 import torch
 from tqdm import tqdm
 import numpy as np
@@ -107,11 +107,13 @@ class SchedulerWrapper:
 		if A.loss_scheduler is _A:B=A.scheduler.set_timesteps(D,**C);A.timesteps=A.scheduler.timesteps;A.init_noise_sigma=A.scheduler.init_noise_sigma;A.order=A.scheduler.order;return B
 		else:B=A.loss_scheduler.set_timesteps(D,**C);A.timesteps=A.loss_scheduler.timesteps;A.init_noise_sigma=A.scheduler.init_noise_sigma;A.order=A.scheduler.order;return B
 	def step(B,model_output,timestep,sample,**F):
+		global image_count
+		global count
 		D=sample;E=model_output;A=timestep; 
-		if timestep==1: print("resetting"); image_count+=1;count=0
+		if timestep==B.timesteps[0]: print("resetting"); image_count+=1;count=0
 		if True:
-			C=B.scheduler.step(E,A,D,**F);A=A.tolist(); print(timestep, B.timesteps)
-			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: 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