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

fix: bug with saving

parent 4446dbff

tune.py

+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))