Commit 5f3ea8e8 authored by Václav Novák's avatar Václav Novák
Browse files

fix: bug with saving

parent 4446dbff
Loading
Loading
Loading
Loading
+1 −1
Original line number Diff line number Diff line
@@ -87,7 +87,7 @@ def objective(trial: Trial) -> int:
        model_name = "model_" + ''.join(random.choice(string.ascii_lowercase) for _ in range(16)) + ".pt"
        
        # Save the model
        torch.save(model.state_dict(), os.path.join(MODEL_DIR, model_name))
        torch.save(model, os.path.join(MODEL_DIR, model_name))
             
        print("Model with AU PRC of {} successfully saved as {}".format(au_prc, model_name))