Run on certain GPUs

[quancs]

🚀 Feature

Can we add a cudavd arg in Trainer to select the gpus to use for the whole program? This feature is important for the user to select the gpus to use (for memory reseans or for usage reasons, or....).

Below is my code for doing this.

  parser = ArgumentParser()
  parser = Trainer.add_argparse_args(parser)
  parser.add_argument('--cudavd', type=str, default='')
  temp_args, _ = parser.parse_known_args()
  if len(temp_args.cudavd) > 0:
      import os
      os.environ["CUDA_VISIBLE_DEVICES"] = temp_args.cudavd

[ananthsub]

Does the trainer gpus flag work for your use case? https://pytorch-lightning.readthedocs.io/en/latest/common/trainer.html