Update src/loss.py
Browse files- src/loss.py +1 -1
src/loss.py
CHANGED
|
@@ -48,7 +48,7 @@ class LearnableScaler(torch.nn.Module):
|
|
| 48 |
super().__init__()
|
| 49 |
self.scale_factors = torch.nn.Parameter(torch.rand(size, size)).to('cuda')
|
| 50 |
def forward(self, timestep):
|
| 51 |
-
return self.scale_factors * torch.log(1 + timestep)
|
| 52 |
|
| 53 |
class LossSchedulerModel(torch.nn.Module):
|
| 54 |
def __init__(A,wx,we, strategy="learnable_scaler"):
|
|
|
|
| 48 |
super().__init__()
|
| 49 |
self.scale_factors = torch.nn.Parameter(torch.rand(size, size)).to('cuda')
|
| 50 |
def forward(self, timestep):
|
| 51 |
+
return self.scale_factors * torch.log(1 + torch.tensor(timestep)).to('cuda')
|
| 52 |
|
| 53 |
class LossSchedulerModel(torch.nn.Module):
|
| 54 |
def __init__(A,wx,we, strategy="learnable_scaler"):
|