I tried it on my side, it works as below.
$ docker run --gpus all -it --rm -v /tmp/.X11-unix:/tmp/.X11-unix -v /home/$user/:/home/$user/ -e DISPLAY=$DISPLAY -w /opt/nvidia/deepstream/deepstream nvcr.io/nvidia/deepstream:5.1-21.02-devel
# apt-get install python3-pip
# pip3 install torch
# root@4661d61ab1f2:~# python3
Python 3.6.9 (default, Jan 26 2021, 15:33:00)
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>>