gaoxiangjun/MPS-NeRF

Train on multi-gpus

Opened this issue · 3 comments

thanks for your great work! I have found that the project seems to support training on multi-gpus. How can I run it with multi-gpus?

As I directly run command

CUDA_VISIBLE_DEVICES=0,1 \
python3 run_nerf_batch.py \
--config configs/canonical_transformer.txt \
--view_num 3 \
--data_root ./data/THuman/nerf_data_/results_gyx_20181013_hyd_1_M \
--expname THuman_3_view

the error occurs

  File "~/projects/MPS-NeRF/lib/skinnning_batch.py", line 400, in forward
    coarse_canonical_pts = self.coarse_deform_target2c(tp_input['params'], tp_input['vertices'], smpl_query_pts)
  File "~/projects/MPS-NeRF/lib/skinnning_batch.py", line 237, in coarse_deform_target2c
    bweights = self.SMPL_NEUTRAL['weights'][vert_ids.squeeze(0)].view(-1,24).cuda()
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cuda:0)

Move index and tensor to the same device:

vert_ids = vert_ids.squeeze(0).cuda()
weights = self.SMPL_NEUTRAL['weights'].cuda()
bweights = weights[vert_ids].view(-1, 24).cuda()

thanks for your reply, i will try it later