FROM hpcaitech/pytorch-cuda:1.10.1-11.3.0

# install colossalai
RUN git clone https://github.com/hpcaitech/ColossalAI.git \
 && cd ./ColossalAI \
 && pip install -v --no-cache-dir .