#!/bin/bash

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