how to save the segmentation model?

[TSbme]

After training, I can't find the best_model.pth. I think the code just calculates the results without saving the model. How to save it?

[KumoLiu]

Hi @tangqishi, you can refer to this tutorial and find how to save model.

if metric > best_metric:
      best_metric = metric
      torch.save(
          model.state_dict(),
          os.path.join(root_dir, "best_metric_model.pth"),
      )
      print("saved new best metric model")

Hope it can help you, thanks!

[TSbme]

Thanks so much！