From ac1d99bd33acb00f8aad3b09e3ce620b03efff44 Mon Sep 17 00:00:00 2001
From: Schneider Leo <leo.schneider@etu.ec-lyon.fr>
Date: Fri, 18 Apr 2025 08:40:54 +0200
Subject: [PATCH] fix : weight Xent loss
---
image_ref/main_ray.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/image_ref/main_ray.py b/image_ref/main_ray.py
index d0ec7c0..9f33926 100644
--- a/image_ref/main_ray.py
+++ b/image_ref/main_ray.py
@@ -53,7 +53,7 @@ def train_model(config,args):
elif config['optimizer']=='SGD' :
optimizer = optim.SGD(model.parameters(), lr=config["lr"], momentum=0.9)
# init training
- n_class = data_train.dataset.classes
+ n_class = len(data_train.dataset.classes)
loss_function = nn.CrossEntropyLoss(weight=torch.Tensor([1/n_class,1-1/n_class]))
# Load existing checkpoint through `get_checkpoint()` API.
--
GitLab