#!/bin/bash

# 检查是否提供了输入参数
if [ -z "$1" ]; then
  echo "please set input image"
  exit 1
fi

# 检查第一个输入参数中是否包含"pytorch"字符串
if [[ "$1" == *"pytorch"* ]]; then
  docker run --rm --platform=linux/amd64 --gpus all $1  python -c \
      "import os; \
      os.system(\"cat /etc/issue\"); \
      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__);
      "
elif [[ "$1" == *"tensorflow"* ]]; then
  docker run --rm --platform=linux/amd64 --gpus all $1  python -c \
      "import os; \
      os.system(\"cat /etc/issue\"); \
      import sys; \
      print(\"python version: \", sys.version); \
      import tensorflow as tf; \
      print(\"tensorflow version: \", tf.__version__); \
      print(\"tensorflow cuda available: \", tf.test.is_gpu_available()); \
      os.system('nvcc -V | tail -n 2')
      "
elif [[ "$1" == *"paddle"* ]]; then
  TARGET_DIR=gpu-base-image-test/paddletest
  docker run --rm --platform=linux/amd64 --gpus all -v ./$TARGET_DIR:/workspace --workdir /workspace $1 python base_test.py

else
  echo "ERROR: no supported test shell"
  exit 1
fi