Update src/loss.py

src/loss.py

@@ -111,7 +111,7 @@ class SchedulerWrapper:
 		A.loss_scheduler=_A
 		A.loss_params_path=loss_params_path
 	def set_timesteps(A,num_inference_steps,**C):
-		D=20
+		D=num_inference_steps
 		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):
@@ -120,7 +120,7 @@ class SchedulerWrapper:
 		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 False:
+		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]=[]
 			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