3_image_test.sh

#!/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/stable-diffusion-v1-4 $1 python infer.py; fi

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