Corrects pytorch_after import path (#407)

Solving issue from MONAI repository: [#8352](Project-MONAI/MONAI#8352) Signed-off-by: R. Garcia-Dias <[email protected]>
Project-MONAI · Mar 2, 2025 · 0ed93da · 0ed93da
1 parent 207cad9
commit 0ed93da
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/auto3dseg/algorithm_templates/dints/scripts/train.py b/auto3dseg/algorithm_templates/dints/scripts/train.py
@@ -27,15 +27,11 @@
 
 import mlflow
 import mlflow.pytorch
+import monai
 import numpy as np
 import torch
 import torch.distributed as dist
 import yaml
-from torch.nn.parallel import DistributedDataParallel
-from torch.utils.tensorboard import SummaryWriter
-from tqdm import tqdm
-
-import monai
 from monai import transforms
 from monai.apps.auto3dseg.auto_runner import logger
 from monai.apps.utils import DEFAULT_FMT
@@ -45,8 +41,11 @@
 from monai.data import DataLoader, partition_dataset
 from monai.inferers import sliding_window_inference
 from monai.metrics import compute_dice
-from monai.networks.utils import pytorch_after
 from monai.utils import RankFilter, set_determinism
+from monai.utils.module import pytorch_after
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
 
 try:
     from apex.contrib.clip_grad import clip_grad_norm_