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 .