fix : weight Xent loss

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.