QuanTrainModule - GPU Support?

[IsidoraR]

🐛 Describe the bug

In the quantize_classification_exmple.py script, there is a warning that says 'quantized inference/test may fail as it is not yet supported in gpu'. However, the Quantization README (https://github.com/TexasInstruments/edgeai-torchvision/blob/master/docs/pixel2pixel/Quantization.md) says that that multi-GPU usage is supported with the QuantTrain Module and QuantTest Module.

Could you please confirm that the QuantTrain and QuantTest modules are supported on a GPU?

Versions

Collecting environment information... PyTorch version: 1.10.0 Is debug build: False CUDA used to build PyTorch: 10.2 ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64) GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0 Clang version: Could not collect CMake version: version 3.16.3 Libc version: glibc-2.31

Python version: 3.8.12 (default, Oct 12 2021, 13:49:34) [GCC 7.5.0] (64-bit runtime) Python platform: Linux-5.13.0-27-generic-x86_64-with-glibc2.17 Is CUDA available: True CUDA runtime version: 10.1.243 GPU models and configuration: Could not collect Nvidia driver version: Could not collect cuDNN version: Could not collect HIP runtime version: N/A MIOpen runtime version: N/A

Versions of relevant libraries: [pip3] numpy==1.21.4 [pip3] onnx-pytorch==0.1.4 [pip3] onnx2pytorch==0.4.1 [pip3] torch==1.10.0+cu111 [pip3] torchaudio==0.10.0 [pip3] torchinfo==1.5.4 [pip3] torchsummary==1.5.1 [pip3] torchvision==0.11.1 [conda] _pytorch_select 0.1 cpu_0
[conda] _tflow_select 2.3.0 mkl
[conda] blas 1.0 mkl
[conda] cudatoolkit 10.2.89 hfd86e86_1
[conda] ffmpeg 4.3 hf484d3e_0 pytorch [conda] mkl 2021.4.0 h06a4308_640
[conda] mkl-service 2.4.0 py38h7f8727e_0
[conda] mkl_fft 1.3.1 py38hd3c417c_0
[conda] mkl_random 1.2.2 py38h51133e4_0
[conda] numpy 1.21.4 pypi_0 pypi [conda] numpy-base 1.21.2 py38h79a1101_0
[conda] onnx-pytorch 0.1.4 pypi_0 pypi [conda] onnx2pytorch 0.4.1 pypi_0 pypi [conda] pytorch 1.10.0 py3.8_cuda10.2_cudnn7.6.5_0 pytorch [conda] pytorch-mutex 1.0 cuda pytorch [conda] tensorflow 2.4.1 mkl_py38hb2083e0_0
[conda] tensorflow-base 2.4.1 mkl_py38h43e0292_0
[conda] torch 1.10.0+cu111 pypi_0 pypi [conda] torchaudio 0.10.0 py38_cu102 pytorch [conda] torchinfo 1.5.4 pypi_0 pypi [conda] torchsummary 1.5.1 pypi_0 pypi [conda] torchvision 0.11.1 pypi_0 pypi