fix(rayTrain): float has no attribute .item() #68

jyaacoub · Dec 14, 2023 · f58d614 · f58d614
1 parent 1bb8d45
commit f58d614
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/rayTrain_Tune.py b/rayTrain_Tune.py
@@ -58,7 +58,7 @@ def train_func(config):
             torch.save(model.state_dict(), checkpoint_path)
             checkpoint = Checkpoint.from_directory(checkpoint_dir)
 
-        ray.train.report({"loss": loss.item()},   checkpoint=checkpoint)
+        ray.train.report({"loss": loss},   checkpoint=checkpoint)
 
 
 if __name__ == "__main__":