File tree 1 file changed +6
-2
lines changed
contrib/MedicalSeg/medicalseg/cvlibs
1 file changed +6
-2
lines changed Original file line number Diff line number Diff line change 21
21
import yaml
22
22
23
23
from medicalseg .cvlibs import manager
24
- from medicalseg .utils import logger
24
+ from medicalseg .utils import logger , get_sys_env
25
25
26
26
# todo: check and edit the unnecessary components
27
27
@@ -318,9 +318,13 @@ def model(self) -> paddle.nn.Layer:
318
318
319
319
if not self ._model :
320
320
self ._model = self ._load_object (model_cfg )
321
- if paddle .get_device () != 'cpu' :
321
+
322
+ env_info = get_sys_env ()
323
+ if paddle .get_device () == 'gpu' and env_info ['Paddle compiled with cuda' ] \
324
+ and env_info ['GPUs used' ] and paddle .distributed .ParallelEnv ().nranks > 1 :
322
325
self ._model = paddle .nn .SyncBatchNorm .convert_sync_batchnorm (
323
326
self ._model )
327
+ logger .info ("Convert bn to sync_bn" )
324
328
325
329
return self ._model
326
330
You can’t perform that action at this time.
0 commit comments