PyTorch GPU Setup for training

I have a GPU setup for Tensorflow 2.6.0.
Cuda v11.2 and Cudnn v8.1.0
Enviroments paths are added.



After that I started learning Unity ML Agents.
I followed the installations instructions.
I didn't get any Cudnn errors.
So I assumed my GPU installation is also done, since it says PyTorch 1.7.1+cu110



But I couldn't train on GPU.
I tried both --torch-device=cuda and --torch-device=cuda:0 parameters.

What would be the proper GPU installation?

 

GPUs only come into use when using visual observations and are leveraging the multiple different CNN architectures (set in the config file for visual observations model). Because most of the observations in ML-Agents are 1D vector observations, there aren't really any gains to get from using a GPU.

On the other hand, if you had a purely vision based agent that used a camera or something similar, the GPU would increase training speed because it's operating on 2D images with 1 - n channels.

You can search this forum for other posts regarding gpu and will find similar responses.