diff --git a/src/ptbench/scripts/saliency_completeness.py b/src/ptbench/scripts/saliency_completeness.py index f17ff648559ff4340fe2aa3179692ac2d6360fd7..2d1de183eecfc85e581b544c4fa948effe26fdf9 100644 --- a/src/ptbench/scripts/saliency_completeness.py +++ b/src/ptbench/scripts/saliency_completeness.py @@ -25,7 +25,7 @@ logger = setup(__name__.split(".")[0], format="%(levelname)s: %(message)s") .. code:: sh - ptbench saliency-completeness -vv pasa tbx11k-v1-healthy-vs-atb --device="cuda" --weight=path/to/model_final.pth --output-folder=path/to/completeness-scores/ + ptbench saliency-completeness -vv pasa tbx11k-v1-healthy-vs-atb --device="cuda" --weight=path/to/model_final.ckpt --output-folder=path/to/completeness-scores/ """, ) diff --git a/src/ptbench/scripts/train.py b/src/ptbench/scripts/train.py index 6eb91fa10080b4cb21f0465b2c7f949cf3ac2bda..c5c2a3ac8ec39017274231ecf2600eb5c043d321 100644 --- a/src/ptbench/scripts/train.py +++ b/src/ptbench/scripts/train.py @@ -257,10 +257,10 @@ def train( """Trains an CNN to perform image classification. Training is performed for a configurable number of epochs, and - generates at least a final_model.pth. It may also generate a number - of intermediate checkpoints. Checkpoints are model files (.pth - files) that are stored during the training and useful to resume the - procedure in case it stops abruptly. + generates at least a final_model.ckpt. It may also generate a + number of intermediate checkpoints. Checkpoints are model files + (.ckpt files) that are stored during the training and useful to + resume the procedure in case it stops abruptly. """ import torch