Add global torch.device variable set based on `--gpus` flag to prevent issues between docker & no-docker device ids

twosixlabs / armory

ARMORY Adversarial Robustness Evaluation Test Bed

MIT License

174 stars 67 forks source link

Add global torch.device variable set based on `--gpus` flag to prevent issues between docker & no-docker device ids #1923

Closed jprokos26 closed 1 year ago

jprokos26 commented 1 year ago

Related to #1902 and #1922

We use torch.device("cuda") a lot in our code: https://github.com/search?q=repo%3Atwosixlabs%2Farmory%20device&type=code

There should instead be a single device variable to be inherited such that we do not need to track the device-id discrepancy between being in a docker container or not. This will also allow for us to set torch.device("cpu") to address #1902.