Update src/loss.py

src/loss.py

@@ -93,7 +93,7 @@ class LossScheduler:
 		global image_count
 		global count
 		A=self;B=A.timesteps.tolist().index(timestep);assert A.t_prev==-1 or B==A.t_prev+1; count+=1
-		if timestep==B.timesteps[0]: print("resetting"); image_count+=1;count=0;
+		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==B.timesteps[-1]: