#!/bin/bash

docker run --rm --platform=linux/amd64 --gpus all $1 python -c \
      'import torch; \
      print("torch version: ", torch.__version__); \
      print("torch cuda available: ", torch.cuda.is_available()); \
      import torchvision; \
      print("torchvision version: ", torchvision.__version__); \
      import torchaudio; \
      print("torchaudio version: ", torchaudio.__version__);
      '