How to use TensorFlow GPU using SSH in Visual Studio Code

  • Install Visual Studio Code
  • Install Remote Development extension pack
  • Install Python Interpreter on remote host (VSC will propose this)
  • Make sure you have correct versions of CUDA Toolkit and CUDA Deep Neural Network – This may require you to change .bashrc and load these modules (perhaps first unloading previous versions of cuda-toolkit and cudnn). The .bashrc file is automatically executed when VSC opens a SSH connection
  • Run a sample Python file containing following content
    import tensorflow as tf
    print(tf.config.list_physical_devices(‘GPU’))”
  • If all was configured well, there should be sensible output from the Python script

Leave a Comment

Scroll to Top