Errors in running infer.py

[xmlyqing00]

Hi,

Nice work! Congratulations. After installed PyTorch and PyTorch3D, I downloaded the checkpoints for DECA and MODNet. But when I type python3 infer.py --deca DECA --rome data. I have the following errors. Do you have any ideas? I ran my code in WSL2, running other PyTorch models are fine.

Start infer!
creating the FLAME Decoder
trained model found. load DECA/data/deca_model.tar
Traceback (most recent call last):
  File "/mnt/c/Users/root/Sources/rome/infer.py", line 218, in <module>
    main(args)
  File "/mnt/c/Users/root/Sources/rome/infer.py", line 195, in main
    infer.run_example()
  File "/mnt/c/Users/root/Sources/rome/infer.py", line 187, in run_example
    out = self.evaluate(source_image, driver_img, crop_center=True)
  File "/home/yq/miniconda3/envs/rome/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/mnt/c/Users/root/Sources/rome/infer.py", line 175, in evaluate
    out = self.model(data_dict,
  File "/home/yq/miniconda3/envs/rome/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/yq/miniconda3/envs/rome/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/mnt/c/Users/root/Sources/rome/src/rome.py", line 243, in forward
    parametric_output = self.parametric_avatar.forward(
  File "/mnt/c/Users/root/Sources/rome/src/parametric_avatar.py", line 565, in forward
    opdict, visdict = self.decode(
  File "/home/yq/miniconda3/envs/rome/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/mnt/c/Users/root/Sources/rome/src/parametric_avatar.py", line 434, in decode
    verts_deforms = self.deform_source_mesh(verts_template, neural_texture, deformer_nets)
  File "/home/yq/miniconda3/envs/rome/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/mnt/c/Users/root/Sources/rome/src/parametric_avatar.py", line 366, in deform_source_mesh
    mlp_input_uv_z = F.grid_sample(uv_deformations_codes, verts_uvs, align_corners=False)[..., 0].permute(0, 2, 1)
  File "/home/yq/miniconda3/envs/rome/lib/python3.9/site-packages/torch/nn/functional.py", line 4223, in grid_sample
    return torch.grid_sampler(input, grid, mode_enum, padding_mode_enum, align_corners)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument grid in method wrapper__grid_sampler_2d)

Thank you for your help!

Best, Yongqing

[xmlyqing00]

Oh, I figure it out. I read the code in rome.py and added an argument --device cuda. The problem was solved. Now I'm looking for where the result is ...