Update src/loss.py

src/loss.py

@@ -118,8 +118,8 @@ class SchedulerWrapper:
 		global image_count
 		global count
 		D=sample;E=model_output;A=timestep; 
-		if timestep==B.timesteps[0]: print("resetting"); image_count+=1;count=0;
-		if timestep==1: print("starting"); image_count+=1;count=0;
+		if timestep==B.timesteps[0]: print("resetting");
+		if timestep: print(timestep); 
 		if B.loss_scheduler is _A:
 			C=B.scheduler.step(E,A,D,**F);A=A.tolist();
 			if A not in B.catch_x:B.catch_x[A]=[];B.catch_e[A]=[];B.catch_x_[A]=[]