#!/bin/bash
source ./venv/bin/activate
GPUS=$1
CONFIG=$2
PORT=1234
if (( $GPUS > 1 )); then
torchrun --nproc_per_node=$GPUS --master_port=$PORT ./src/train.py -opt "$CONFIG" --launcher=pytorch
else
python3 ./src/train.py -opt "$CONFIG"
fi
deactivate