You have already included machine learning packages. Serious machine learning models are computed using an Nvidia GPU. The container should include Nvidia drivers and probably CUDNN. I've included some python code to verify that Tensorflow and Torch have access to a GPU.
Python code to test if Tensorflow has access to an Nvidia GPU:
# TestTensorflowForGpu.py
print("Requirement 1: Install Nvidia CUDA Tookit. See: https://towardsdatascience.com/installing-tensorflow-with-cuda-cudnn-and-gpu-support-on-windows-10-60693e46e781")
print("Requirement 2: Install CUDNN. See: https://medium.com/analytics-vidhya/installing-cudnn-for-gpu-support-with-tensorflow-on-windows-10-aff10c6c9929")
print("If cudart64_101.dll is missing. see: https://stackoverflow.com/questions/59823283/could-not-load-dynamic-library-cudart64-101-dll-on-tensorflow-cpu-only-install")
print("If missing Dlls are reported: see: https://github.com/tensorflow/tensorflow/issues/44291")
import tensorflow as tf
tf.config.list_physical_devices('GPU')
Python code to test if Torch has access to an Nvidia GPU:
# TestTorchForGpu.py
print("Requirement 1: Install Nvidia CUDA Tookit. See: https://towardsdatascience.com/installing-tensorflow-with-cuda-cudnn-and-gpu-support-on-windows-10-60693e46e781")
print("Requirement 2: Install CUDNN. See: https://medium.com/analytics-vidhya/installing-cudnn-for-gpu-support-with-tensorflow-on-windows-10-aff10c6c9929")
print("If cudart64_101.dll is missing. see: https://stackoverflow.com/questions/59823283/could-not-load-dynamic-library-cudart64-101-dll-on-tensorflow-cpu-only-install")
print("If missing Dlls are reported: see: https://github.com/tensorflow/tensorflow/issues/44291")
# if you get 'AssertionError: Torch not compiled with CUDA enabled'
# see: https://github.com/pytorch/pytorch/issues/30664
# 1) pip uninstall torch
# 2) pip cache purge
# 3) pip install torch -f https://download.pytorch.org/whl/torch_stable.html
import torch
print(f'torch.cuda.device.count: {torch.cuda.device_count()}')
print(f'torch.cuda.get_device_name: {torch.cuda.get_device_name()}')
You have already included machine learning packages. Serious machine learning models are computed using an Nvidia GPU. The container should include Nvidia drivers and probably CUDNN. I've included some python code to verify that Tensorflow and Torch have access to a GPU.
Python code to test if Tensorflow has access to an Nvidia GPU:
Python code to test if Torch has access to an Nvidia GPU: