droid-dataset / droid_policy_learning

DROID Policy Learning and Evaluation
MIT License
112 stars 12 forks source link

Multi-GPU training #19

Open Lizhuoling opened 4 weeks ago

Lizhuoling commented 4 weeks ago

Does the code support distributed training using multiple GPUs? Or the training only needs 1 GPU so multi-gpu training is not needed?

kpertsch commented 3 weeks ago

Our current training is all run on a single GPU (A100) so the code does not support multi-GPU training at the moment. It shouldn't be hard to add though by using PyTorch DDP.