Skip to content

Commit

Permalink
Merge pull request #66 from zooniverse/zoobot-upgrade-testing
Browse files Browse the repository at this point in the history
add gpus all command to container_run_options
  • Loading branch information
Tooyosi authored Dec 19, 2024
2 parents ad2264a + 0546f7c commit 20cfbdd
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions bajor/batch/train_finetuning.py
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,7 @@ def create_job_tasks(job_id, task_id=1, run_opts=''):
command = (
'/bin/bash -c "'
'set -ex; '
'nvidia-smi; '
'nvidia-smi; || echo \\"nvidia-smi command failed\\"; '
'python -c \\"import torch; print(torch.cuda.is_available()); print(torch.cuda.device_count())\\"; '
f'{wait_for_preparation_task_completion}; '
f'{setup_pytorch_kernel_cache_env_var}; '
Expand All @@ -251,7 +251,7 @@ def create_job_tasks(job_id, task_id=1, run_opts=''):
container_settings=batchmodels.TaskContainerSettings(
image_name=os.getenv('CONTAINER_IMAGE_NAME'),
working_directory='taskWorkingDirectory',
container_run_options='--ipc=host'
container_run_options='--ipc=host --gpus all'
),
user_identity = batchmodels.UserIdentity(
auto_user=batchmodels.AutoUserSpecification(
Expand Down

0 comments on commit 20cfbdd

Please sign in to comment.