Update src/loss.py
Browse files- src/loss.py +1 -1
src/loss.py
CHANGED
|
@@ -151,7 +151,7 @@ class LossScheduler:
|
|
| 151 |
A.optimizer.step()
|
| 152 |
if B+1==len(A.timesteps):
|
| 153 |
A.xT=_A;A.e_prev=[]
|
| 154 |
-
A.t_prev=-1
|
| 155 |
A.save("/home/mbhat/weights.pth")
|
| 156 |
else:A.t_prev=B
|
| 157 |
return C,
|
|
|
|
| 151 |
A.optimizer.step()
|
| 152 |
if B+1==len(A.timesteps):
|
| 153 |
A.xT=_A;A.e_prev=[]
|
| 154 |
+
A.t_prev=-1; print(model_output.shape, model_output)
|
| 155 |
A.save("/home/mbhat/weights.pth")
|
| 156 |
else:A.t_prev=B
|
| 157 |
return C,
|