diff --git a/main_ray_tune.py b/main_ray_tune.py index 85542043af7240001774ed7e684db8b718242c85..00d12bc559da4f43ebd0f68290c474b75486ab46 100644 --- a/main_ray_tune.py +++ b/main_ray_tune.py @@ -158,7 +158,7 @@ def train_model(config, args): loss_rt = criterion_rt(rt, pred_rt) loss_int = criterion_intensity(intensity, pred_int) loss = loss_rt + loss_int - val_loss += loss.item().numpy() + val_loss += loss.item() val_steps += 1 # Here we save a checkpoint. It is automatically registered with