about multi GPU in tran.py

askerlee / segtran

Medical Image Segmentation using Squeeze-and-Expansion Transformers

214 stars 50 forks source link

about multi GPU in tran.py #3

Closed BossZard closed 2 years ago

BossZard commented 3 years ago

Nice job! but i want to use multi for training , Can you tell me which parameter can be modified to use multi GPU? THX!

askerlee commented 3 years ago

You can try the command below: python3 -m torch.distributed.launch --nproc_per_node=2 train2d.py .... But the multi-GPU code hasn't been tested for a while. In case there's any error please don't hesitate to inform me 😃