mesh/scripts/train.sh

#!/bin/bash

export CUDA_VISIBLE_DEVICES=4,5,6,7

torchrun \
    --nproc_per_node=4 \
    --master_port=29889 \
    train.py \
    --batch_size 2 \
    --epochs 100 \
    --save_every_k_epochs 1