Wrong pytorch-cuda version installation with pip

[mmeendez8]

Describe the bug
It seems that this line:
https://github.com/grimoire/mmdetection-to-tensorrt/blob/master/docker/Dockerfile#L21

It is not working as intended since it is installing pytorch with cuda=10.2.

I am working with a RTX3060 and I am getting the following error:

NVIDIA GeForce RTX 3060 Laptop GPU with CUDA capability sm_86 is not compatible with the current PyTorch installation.
The current PyTorch install supports CUDA capabilities sm_37 sm_50 sm_60 sm_70 sm_75.
If you want to use the NVIDIA GeForce RTX 3060 Laptop GPU GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/

This is related with a wrong installation of pytorch (more info in this thread)

Basically we are using a base tensorrt image that contains cuda 11.1,1:
https://docs.nvidia.com/deeplearning/tensorrt/container-release-notes/rel_20-12.html#rel_20-12

But it seems pytorch installation is installing cuda 10.2 binaries. You can check this in the official documentation:
https://pytorch.org/get-started/locally/ where basically if you choose pip and cuda 10.2 they return the following command:

pip3 install torch torchvision torchaudio

We should modify torch installation to add cuda specific version.

I am willing to send a PR that modifies this!

[mmeendez8]

Actually note that it does not seem to exist binaries for cuda 11.1 for pytorch<1.8
See full list here:
https://pytorch.org/get-started/previous-versions/

[grimoire]

Thank you! PRs are welcome!
You can update the base image if there the cuda version and pytorch can not match.