Update src/loss.py

src/loss.py

@@ -5,7 +5,7 @@ import output
 class LossSchedulerModel(torch.nn.Module):
 	def __init__(A,wx,we):super(LossSchedulerModel,A).__init__();assert len(wx.shape)==1 and len(we.shape)==2;B=wx.shape[0];assert B==we.shape[0]and B==we.shape[1];A.register_parameter('wx',torch.nn.Parameter(wx));A.register_parameter('we',torch.nn.Parameter(we))
 	def forward(A,t,xT,e_prev):
-		B=e_prev;assert t-len(B)==0;C=xT*A.wx[t]
+		B=e_prev;assert t-len(B)+1==0;C=xT*A.wx[t]
 		for(D,E)in zip(B,A.we[t]):C+=D*E
 		return C.to(xT.dtype)
 class LossScheduler: