Update src/loss.py

src/loss.py

@@ -31,7 +31,7 @@ class ExponentialScalingMatrix(torch.nn.Module):
 		super().__init__()
 		self.base = torch.nn.Parameter(torch.ones(size, size)).to('cuda')
 	def forward(self, timestep):
-		return self.base * torch.exp(-timestep / 10.0).to('cuda')
+		return self.base * torch.exp(-torch.tensor(timestep) / 10.0).to('cuda')
 
 class HierarchicalMatrix(torch.nn.Module):
 	def __init__(self, size):