FROM python:3.9-slim-buster
# works with 3.8
USER root
WORKDIR /root
RUN apt-get update && apt-get install git build-essential --yes
RUN python3 -m pip install --upgrade pip
RUN python3 -m pip install --pre torch==1.8.0.dev20210102+cpu  -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
RUN python3 -m pip install Pillow
RUN git clone https://github.com/pytorch/vision/
ENV PYTHONPATH=./vision
COPY test-backbone.py /
RUN python /test-backbone.py

from torchvision.models import resnet50
import torch


if __name__ == "__main__":
    backbone = resnet50(pretrained=True)

    img = torch.randn((2, 3, 256, 288))
    features = backbone(img)
    loss = features.sum()
    loss.backward()
    print("exiting")

Step 11/11 : RUN python /test-backbone.py
 ---> Running in 31a88ed05b3d
Downloading: "https://download.pytorch.org/models/resnet50-19c8e357.pth" to /root/.cache/torch/hub/checkpoints/resnet50-19c8e357.pth
100.0%
exiting
Segmentation fault (core dumped)
The command '/bin/sh -c python /test-backbone.py' returned a non-zero code: 139