@@ -1,4 +1,5 @@
-FROM nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04 AS builder
+# FROM nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04 AS builder
+FROM pytorch/pytorch:1.7.1-cuda11.0-cudnn8-devel AS builder
RUN sed -i 's#archive.ubuntu.com#mirrors.aliyun.com#g' /etc/apt/sources.list \
&& sed -i 's#security.ubuntu.com#mirrors.aliyun.com#g' /etc/apt/sources.list
