Update src/loss.py

src/loss.py

@@ -96,7 +96,7 @@ class LossScheduler:
 		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]:
+		if timestep==A.timesteps[-1]:
 			with torch.enable_grad():
 				orig_output = torch.load(f"/home/mbhat/edge-maxxing/miner/miner/latents/latent_oig_{image_count}_20.pth")
 				A.optimizer.zero_grad()
@@ -107,7 +107,7 @@ class LossScheduler:
 		if B+1==len(A.timesteps):
 			A.xT=_A;A.e_prev=[];
 			A.t_prev=-1
-			A.save("/home/mbhat/weights_latest.pth")    
+			A.save(f"/home/mbhat/weights/weights_latest_{int(timestep)}.pth")    
 		else:A.t_prev=B
 		return C,
 	def step_other(self,model_output,timestep,sample,*D,**E):