From 7bd07e8a345bb508486b27393fdf270e7fd4f65b Mon Sep 17 00:00:00 2001 From: valhassan Date: Mon, 27 May 2024 16:01:05 -0400 Subject: [PATCH] Add logging statement and update model initialization parameters --- train_segmentation.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/train_segmentation.py b/train_segmentation.py index cf2a7ca0..3376e68c 100644 --- a/train_segmentation.py +++ b/train_segmentation.py @@ -797,13 +797,15 @@ def train(cfg: DictConfig) -> None: # Script model if scriptmodel: + logging.info(f'\nScripting model...') model_to_script = ScriptModel(model, device=device, + num_classes=num_classes, input_shape=(1, num_bands, patches_size, patches_size), mean=mean, std=std, - min=scale[0], - max=scale[1]) + scaled_min=scale[0], + scaled_max=scale[1]) scripted_model = torch.jit.script(model_to_script) scripted_model.save(output_path.joinpath('scripted_model.pt'))