Local trainning with GPU

Question/Issue:
I need some help on local training with GPU

Project ID:

Context/Use case:
So I have follow this to use docker create an image to train model locally.
Once select “edit block locally” option in Classification block, it will download a folder with a train.py.

After run docker with command: “docker run --gpus all --rm -v “%cd%”:/scripts custom-block-650889 --data-directory data --out-directory out” (as suggest in readme.md)
It starts working, but it only print out the gpu info instead of using it to train.


So I wonder what should I do to train with GPU?