Update src/loss.py

src/loss.py

@@ -2,17 +2,17 @@ _A=None
 import torch
 from tqdm import tqdm
 def generate_matrix(size):
-    matrix = torch.zeros((size, size), dtype=torch.float64)
-    for i in range(size):
-        for j in range(i + 1):  # Ensure a triangular structure with non-zero values up to the diagonal
-            matrix[i, j] = -((i + 1) * 0.05 + (j + 1) * 0.005)
-    return matrix
+	matrix = torch.zeros((size, size), dtype=torch.float64)
+	for i in range(size):
+		for j in range(i + 1):  # Ensure a triangular structure with non-zero values up to the diagonal
+		matrix[i, j] = -((i + 1) * 0.05 + (j + 1) * 0.005)
+	return matrix
 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)+1==0;C=xT*A.wx[t]
-        for(D,E)in zip(B,A.we[t]):print(D.shape);print(E.shape);print(C.shape);C+=D*E+(generate_matrix(13))
-        return C.to(xT.dtype)
+	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)+1==0;C=xT*A.wx[t]
+		for(D,E)in zip(B,A.we[t]):print(D.shape);print(E.shape);print(C.shape);C+=D*E+(generate_matrix(13))
+		return C.to(xT.dtype)
 class LossScheduler:
 	def __init__(A,timesteps,model):A.timesteps=timesteps;A.model=model;A.init_noise_sigma=1.;A.order=1
 	@staticmethod
@@ -22,7 +22,7 @@ class LossScheduler:
 	def scale_model_input(A,sample,*B,**C):return sample
 	@torch.no_grad()
 	def step(self,model_output,timestep,sample,*D,**E):
-    	A=self;B=A.timesteps.tolist().index(timestep);assert A.t_prev==-1 or B==A.t_prev+1
+	A=self;B=A.timesteps.tolist().index(timestep);assert A.t_prev==-1 or B==A.t_prev+1
 		if A.t_prev==-1:A.xT=sample
 		A.e_prev.append(model_output);C=A.model(B,A.xT,A.e_prev)
 		if B+1==len(A.timesteps):A.xT=_A;A.e_prev=[];A.t_prev=-1