#!/bin/bash
TARGET_DIR=gpu-base-image-test
# 检查是否提供了输入参数
if [ -z "$1" ]; then
  echo "please set input image"
  exit 1
fi

if [[ "$1" == *"pytorch"* ]]; then \
  docker run --rm --platform=linux/amd64 --gpus all -v ./$TARGET_DIR:/workspace --workdir /workspace/pytorch/gpt2 $1 python infer.py; fi

if [[ "$1" == *"tensorflow"* ]]; then \
  docker run --rm --platform=linux/amd64 --gpus all -v ./$TARGET_DIR:/workspace --workdir /workspace/tensorflow/bert $1 python infer.py; fi
if [[ "$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 text.py; fi