MPE 2.5

martinnovaak · martinnovaak · commit 24098d7b3e86 · 2024-08-07T23:05:45.000+02:00
diff --git a/trainer/train.py b/trainer/train.py
@@ -64,7 +64,7 @@ def train(model: torch.nn.Module, optimizer: torch.optim.Optimizer, dataloader:
         optimizer.zero_grad()
         prediction = model(batch)
 
-        loss = torch.mean((prediction - batch.target) ** 2)
+        loss = torch.mean(torch.abs(prediction - batch.target) ** 2.5)
         loss.backward()
         optimizer.step()
         model.clamp_weights()
