diff --git a/PhysicsTools/TensorFlow/BuildFile.xml b/PhysicsTools/TensorFlow/BuildFile.xml
index 90d48104241aa..0fa15c43a0f7b 100644
--- a/PhysicsTools/TensorFlow/BuildFile.xml
+++ b/PhysicsTools/TensorFlow/BuildFile.xml
@@ -3,6 +3,7 @@
 <use name="FWCore/Framework"/>
 <use name="FWCore/MessageLogger"/>
 <use name="FWCore/Utilities"/>
+<use name="FWCore/ServiceRegistry"/>
 <export>
   <lib name="1"/>
 </export>
diff --git a/PhysicsTools/TensorFlow/interface/TensorFlow.h b/PhysicsTools/TensorFlow/interface/TensorFlow.h
index edb6c86f9bc4f..371c49d4aad04 100644
--- a/PhysicsTools/TensorFlow/interface/TensorFlow.h
+++ b/PhysicsTools/TensorFlow/interface/TensorFlow.h
@@ -25,6 +25,8 @@
 
 namespace tensorflow {
 
+  enum class Backend { cpu, cuda, rocm, intel, best };
+
   typedef std::pair<std::string, Tensor> NamedTensor;
   typedef std::vector<NamedTensor> NamedTensorList;
 
@@ -39,6 +41,10 @@ namespace tensorflow {
   // since the threading configuration is done per run() call as of 2.1
   void setThreading(SessionOptions& sessionOptions, int nThreads, const std::string& singleThreadPool);
 
+  // Set the backend option cpu/cuda
+  // The gpu memory is set to "allow_growth" to avoid TF getting all the CUDA memory at once.
+  void setBackend(SessionOptions& sessionOptions, Backend backend = Backend::cpu);
+
   // loads a meta graph definition saved at exportDir using the SavedModel interface for a tag and
   // predefined sessionOptions
   // transfers ownership
@@ -52,11 +58,13 @@ namespace tensorflow {
   // transfers ownership
   MetaGraphDef* loadMetaGraphDef(const std::string& exportDir,
                                  const std::string& tag = kSavedModelTagServe,
+                                 Backend backend = Backend::cpu,
                                  int nThreads = 1);
 
   // deprecated in favor of loadMetaGraphDef
   MetaGraphDef* loadMetaGraph(const std::string& exportDir,
                               const std::string& tag = kSavedModelTagServe,
+                              Backend backend = Backend::cpu,
                               int nThreads = 1);
 
   // loads a graph definition saved as a protobuf file at pbFile
@@ -67,9 +75,9 @@ namespace tensorflow {
   // transfers ownership
   Session* createSession(SessionOptions& sessionOptions);
 
-  // return a new, empty session with nThreads
+  // return a new, empty session with nThreads and selected backend
   // transfers ownership
-  Session* createSession(int nThreads = 1);
+  Session* createSession(Backend backend = Backend::cpu, int nThreads = 1);
 
   // return a new session that will contain an already loaded meta graph whose exportDir must be
   // given in order to load and initialize the variables, sessionOptions are predefined
@@ -83,7 +91,10 @@ namespace tensorflow {
   // in order to load and initialize the variables, threading options are inferred from nThreads
   // an error is thrown when metaGraphDef is a nullptr or when the graph has no nodes
   // transfers ownership
-  Session* createSession(const MetaGraphDef* metaGraphDef, const std::string& exportDir, int nThreads = 1);
+  Session* createSession(const MetaGraphDef* metaGraphDef,
+                         const std::string& exportDir,
+                         Backend backend = Backend::cpu,
+                         int nThreads = 1);
 
   // return a new session that will contain an already loaded graph def, sessionOptions are predefined
   // an error is thrown when graphDef is a nullptr or when the graph has no nodes
@@ -94,7 +105,7 @@ namespace tensorflow {
   // inferred from nThreads
   // an error is thrown when graphDef is a nullptr or when the graph has no nodes
   // transfers ownership
-  Session* createSession(const GraphDef* graphDef, int nThreads = 1);
+  Session* createSession(const GraphDef* graphDef, Backend backend = Backend::cpu, int nThreads = 1);
 
   // closes a session, calls its destructor, resets the pointer, and returns true on success
   bool closeSession(Session*& session);
diff --git a/PhysicsTools/TensorFlow/plugins/TfGraphDefProducer.cc b/PhysicsTools/TensorFlow/plugins/TfGraphDefProducer.cc
index a749a9f70bfde..2468a74737a6f 100644
--- a/PhysicsTools/TensorFlow/plugins/TfGraphDefProducer.cc
+++ b/PhysicsTools/TensorFlow/plugins/TfGraphDefProducer.cc
@@ -48,7 +48,7 @@ TfGraphDefProducer::TfGraphDefProducer(const edm::ParameterSet& iConfig)
 // ------------ method called to produce the data  ------------
 TfGraphDefProducer::ReturnType TfGraphDefProducer::produce(const TfGraphRecord& iRecord) {
   auto* graph = tensorflow::loadGraphDef(filename_);
-  return std::make_unique<TfGraphDefWrapper>(tensorflow::createSession(graph, 1), graph);
+  return std::make_unique<TfGraphDefWrapper>(tensorflow::createSession(graph), graph);
 }
 
 void TfGraphDefProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
diff --git a/PhysicsTools/TensorFlow/src/TensorFlow.cc b/PhysicsTools/TensorFlow/src/TensorFlow.cc
index 37d3c5183a243..3dffc4c90ab39 100644
--- a/PhysicsTools/TensorFlow/src/TensorFlow.cc
+++ b/PhysicsTools/TensorFlow/src/TensorFlow.cc
@@ -6,8 +6,9 @@
  */
 
 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
-
 #include "FWCore/MessageLogger/interface/MessageLogger.h"
+#include "FWCore/ServiceRegistry/interface/Service.h"
+#include "FWCore/Utilities/interface/ResourceInformation.h"
 
 namespace tensorflow {
 
@@ -25,6 +26,65 @@ namespace tensorflow {
     setThreading(sessionOptions, nThreads);
   }
 
+  void setBackend(SessionOptions& sessionOptions, Backend backend) {
+    /*
+     * The TensorFlow backend configures the available devices using options provided in the sessionOptions proto.
+     * // Options from https://github.com/tensorflow/tensorflow/blob/c53dab9fbc9de4ea8b1df59041a5ffd3987328c3/tensorflow/core/protobuf/config.proto
+     *
+     * If the device_count["GPU"] = 0 GPUs are not used. 
+     * The visible_device_list configuration is used to map the `visible` devices (from CUDA_VISIBLE_DEVICES) to `virtual` devices.
+     * If Backend::cpu is request, the GPU device is disallowed by device_count configuration.
+     * If Backend::cuda is request:
+     *  - if ResourceInformation shows an available Nvidia GPU device:
+     *     the device is used with memory_growth configuration (not allocating all cuda memory at once).
+     *  - if no device is present: an exception is raised.
+     */
+
+    edm::Service<edm::ResourceInformation> ri;
+    if (backend == Backend::cpu) {
+      // disable GPU usage
+      (*sessionOptions.config.mutable_device_count())["GPU"] = 0;
+      sessionOptions.config.mutable_gpu_options()->set_visible_device_list("");
+    }
+    // NVidia GPU
+    else if (backend == Backend::cuda) {
+      if (not ri->nvidiaDriverVersion().empty()) {
+        // Take only the first GPU in the CUDA_VISIBLE_DEVICE list
+        (*sessionOptions.config.mutable_device_count())["GPU"] = 1;
+        sessionOptions.config.mutable_gpu_options()->set_visible_device_list("0");
+        // Do not allocate all the memory on the GPU at the beginning.
+        sessionOptions.config.mutable_gpu_options()->set_allow_growth(true);
+      } else {
+        edm::Exception ex(edm::errors::UnavailableAccelerator);
+        ex << "Cuda backend requested, but no NVIDIA GPU available in the job";
+        ex.addContext("Calling tensorflow::setBackend()");
+        throw ex;
+      }
+    }
+    // ROCm and Intel GPU are still not supported
+    else if ((backend == Backend::rocm) || (backend == Backend::intel)) {
+      edm::Exception ex(edm::errors::UnavailableAccelerator);
+      ex << "ROCm/Intel GPU backend requested, but TF is not compiled yet for this platform";
+      ex.addContext("Calling tensorflow::setBackend()");
+      throw ex;
+    }
+    // Get NVidia GPU if possible or fallback to CPU
+    else if (backend == Backend::best) {
+      // Check if a Nvidia GPU is availabl
+      if (not ri->nvidiaDriverVersion().empty()) {
+        // Take only the first GPU in the CUDA_VISIBLE_DEVICE list
+        (*sessionOptions.config.mutable_device_count())["GPU"] = 1;
+        sessionOptions.config.mutable_gpu_options()->set_visible_device_list("0");
+        // Do not allocate all the memory on the GPU at the beginning.
+        sessionOptions.config.mutable_gpu_options()->set_allow_growth(true);
+      } else {
+        // Just CPU support
+        (*sessionOptions.config.mutable_device_count())["GPU"] = 0;
+        sessionOptions.config.mutable_gpu_options()->set_visible_device_list("");
+      }
+    }
+  }
+
   MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag, SessionOptions& sessionOptions) {
     // objects to load the graph
     Status status;
@@ -49,19 +109,20 @@ namespace tensorflow {
     return loadMetaGraphDef(exportDir, tag, sessionOptions);
   }
 
-  MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag, int nThreads) {
+  MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag, Backend backend, int nThreads) {
     // create session options and set thread options
     SessionOptions sessionOptions;
     setThreading(sessionOptions, nThreads);
+    setBackend(sessionOptions, backend);
 
     return loadMetaGraphDef(exportDir, tag, sessionOptions);
   }
 
-  MetaGraphDef* loadMetaGraph(const std::string& exportDir, const std::string& tag, int nThreads) {
+  MetaGraphDef* loadMetaGraph(const std::string& exportDir, const std::string& tag, Backend backend, int nThreads) {
     edm::LogInfo("PhysicsTools/TensorFlow")
         << "tensorflow::loadMetaGraph() is deprecated, use tensorflow::loadMetaGraphDef() instead";
 
-    return loadMetaGraphDef(exportDir, tag, nThreads);
+    return loadMetaGraphDef(exportDir, tag, backend, nThreads);
   }
 
   GraphDef* loadGraphDef(const std::string& pbFile) {
@@ -95,10 +156,11 @@ namespace tensorflow {
     return session;
   }
 
-  Session* createSession(int nThreads) {
+  Session* createSession(Backend backend, int nThreads) {
     // create session options and set thread options
     SessionOptions sessionOptions;
     setThreading(sessionOptions, nThreads);
+    setBackend(sessionOptions, backend);
 
     return createSession(sessionOptions);
   }
@@ -152,10 +214,11 @@ namespace tensorflow {
     return session;
   }
 
-  Session* createSession(const MetaGraphDef* metaGraphDef, const std::string& exportDir, int nThreads) {
+  Session* createSession(const MetaGraphDef* metaGraphDef, const std::string& exportDir, Backend backend, int nThreads) {
     // create session options and set thread options
     SessionOptions sessionOptions;
     setThreading(sessionOptions, nThreads);
+    setBackend(sessionOptions, backend);
 
     return createSession(metaGraphDef, exportDir, sessionOptions);
   }
@@ -186,10 +249,11 @@ namespace tensorflow {
     return session;
   }
 
-  Session* createSession(const GraphDef* graphDef, int nThreads) {
+  Session* createSession(const GraphDef* graphDef, Backend backend, int nThreads) {
     // create session options and set thread options
     SessionOptions sessionOptions;
     setThreading(sessionOptions, nThreads);
+    setBackend(sessionOptions, backend);
 
     return createSession(graphDef, sessionOptions);
   }
diff --git a/PhysicsTools/TensorFlow/test/BuildFile.xml b/PhysicsTools/TensorFlow/test/BuildFile.xml
index b5ed26c3fedf5..1270225afc9f0 100644
--- a/PhysicsTools/TensorFlow/test/BuildFile.xml
+++ b/PhysicsTools/TensorFlow/test/BuildFile.xml
@@ -2,38 +2,166 @@
   <use name="tensorflow-cc"/>
   <use name="boost_filesystem"/>
   <use name="cppunit"/>
+  <use name="PhysicsTools/TensorFlow"/>
 </bin>
 
+<iftool name="cuda">
+  <bin name="testTFHelloWorldCUDA" file="testRunner.cpp,testHelloWorldCUDA.cc">
+    <use name="tensorflow-cc"/>
+    <use name="boost_filesystem"/>
+    <use name="cppunit"/>
+    <use name="PhysicsTools/TensorFlow"/>
+    <use name="HeterogeneousCore/CUDAServices"/>
+    <use name="HeterogeneousCore/CUDAUtilities"/>
+    <use name="catch2"/>
+    <use name="FWCore/ParameterSet"/>
+    <use name="FWCore/ParameterSetReader"/>
+    <use name="FWCore/PluginManager"/>
+    <use name="FWCore/ServiceRegistry"/>
+    <use name="FWCore/Utilities"/>
+    <use name="cuda"/>
+  </bin>
+</iftool>
+
+
 <bin name="testTFMetaGraphLoading" file="testRunner.cpp,testMetaGraphLoading.cc">
   <use name="boost_filesystem"/>
   <use name="cppunit"/>
   <use name="PhysicsTools/TensorFlow"/>
 </bin>
 
+<iftool name="cuda">
+<bin name="testTFMetaGraphLoadingCUDA" file="testRunner.cpp,testMetaGraphLoadingCUDA.cc">
+  <use name="boost_filesystem"/>
+  <use name="cppunit"/>
+  <use name="PhysicsTools/TensorFlow"/>
+  <use name="HeterogeneousCore/CUDAServices"/>
+  <use name="HeterogeneousCore/CUDAUtilities"/>
+  <use name="catch2"/>
+  <use name="FWCore/ParameterSet"/>
+  <use name="FWCore/ParameterSetReader"/>
+  <use name="FWCore/PluginManager"/>
+  <use name="FWCore/ServiceRegistry"/>
+  <use name="FWCore/Utilities"/>
+  <use name="cuda"/>
+</bin>
+</iftool>
+
 <bin name="testTFGraphLoading" file="testRunner.cpp,testGraphLoading.cc">
   <use name="boost_filesystem"/>
   <use name="cppunit"/>
   <use name="PhysicsTools/TensorFlow"/>
 </bin>
 
+<iftool name="cuda">
+<bin name="testTFGraphLoadingCUDA" file="testRunner.cpp,testGraphLoadingCUDA.cc">
+  <use name="boost_filesystem"/>
+  <use name="cppunit"/>
+  <use name="PhysicsTools/TensorFlow"/>
+  <use name="HeterogeneousCore/CUDAServices"/>
+  <use name="HeterogeneousCore/CUDAUtilities"/>
+  <use name="catch2"/>
+  <use name="FWCore/ParameterSet"/>
+  <use name="FWCore/ParameterSetReader"/>
+  <use name="FWCore/PluginManager"/>
+  <use name="FWCore/ServiceRegistry"/>
+  <use name="FWCore/Utilities"/>
+  <use name="cuda"/>
+</bin>
+</iftool>
+
 <bin name="testTFConstSession" file="testRunner.cpp,testConstSession.cc">
   <use name="boost_filesystem"/>
   <use name="cppunit"/>
   <use name="PhysicsTools/TensorFlow"/>
 </bin>
 
+<iftool name="cuda">
+<bin name="testTFConstSessionCUDA" file="testRunner.cpp,testConstSessionCUDA.cc">
+  <use name="boost_filesystem"/>
+  <use name="cppunit"/>
+  <use name="PhysicsTools/TensorFlow"/>
+  <use name="HeterogeneousCore/CUDAServices"/>
+  <use name="HeterogeneousCore/CUDAUtilities"/>
+  <use name="catch2"/>
+  <use name="FWCore/ParameterSet"/>
+  <use name="FWCore/ParameterSetReader"/>
+  <use name="FWCore/PluginManager"/>
+  <use name="FWCore/ServiceRegistry"/>
+  <use name="FWCore/Utilities"/>
+  <use name="cuda"/>
+</bin>
+</iftool>
+
 <bin name="testTFSessionCache" file="testRunner.cpp,testSessionCache.cc">
   <use name="boost_filesystem"/>
   <use name="cppunit"/>
   <use name="PhysicsTools/TensorFlow"/>
 </bin>
 
+<iftool name="cuda">
+<bin name="testTFSessionCacheCUDA" file="testRunner.cpp,testSessionCacheCUDA.cc">
+  <use name="boost_filesystem"/>
+  <use name="cppunit"/>
+  <use name="PhysicsTools/TensorFlow"/>
+  <use name="HeterogeneousCore/CUDAServices"/>
+  <use name="HeterogeneousCore/CUDAUtilities"/>
+  <use name="catch2"/>
+  <use name="FWCore/ParameterSet"/>
+  <use name="FWCore/ParameterSetReader"/>
+  <use name="FWCore/PluginManager"/>
+  <use name="FWCore/ServiceRegistry"/>
+  <use name="FWCore/Utilities"/>
+  <use name="cuda"/>
+</bin>
+</iftool>
+
 <bin name="testTFThreadPools" file="testRunner.cpp,testThreadPools.cc">
   <use name="boost_filesystem"/>
   <use name="cppunit"/>
   <use name="PhysicsTools/TensorFlow"/>
 </bin>
 
+<iftool name="cuda">
+<bin name="testTFThreadPoolsCUDA" file="testRunner.cpp,testThreadPoolsCUDA.cc">
+  <use name="boost_filesystem"/>
+  <use name="cppunit"/>
+  <use name="PhysicsTools/TensorFlow"/>
+  <use name="HeterogeneousCore/CUDAServices"/>
+  <use name="HeterogeneousCore/CUDAUtilities"/>
+  <use name="catch2"/>
+  <use name="FWCore/ParameterSet"/>
+  <use name="FWCore/ParameterSetReader"/>
+  <use name="FWCore/PluginManager"/>
+  <use name="FWCore/ServiceRegistry"/>
+  <use name="FWCore/Utilities"/>
+  <use name="cuda"/>
+</bin>
+</iftool>
+
+<bin name="testTFVisibleDevices" file="testRunner.cpp,testVisibleDevices.cc">
+  <use name="boost_filesystem"/>
+  <use name="cppunit"/>
+  <use name="PhysicsTools/TensorFlow"/>
+</bin>
+
+
+<iftool name="cuda">
+<bin name="testTFVisibleDevicesCUDA" file="testRunner.cpp,testVisibleDevicesCUDA.cc">
+  <use name="boost_filesystem"/>
+  <use name="cppunit"/>
+  <use name="PhysicsTools/TensorFlow"/>
+  <use name="HeterogeneousCore/CUDAServices"/>
+  <use name="catch2"/>
+  <use name="FWCore/ParameterSet"/>
+  <use name="FWCore/ParameterSetReader"/>
+  <use name="FWCore/PluginManager"/>
+  <use name="FWCore/ServiceRegistry"/>
+  <use name="FWCore/Utilities"/>
+  <use name="cuda"/>
+</bin>
+</iftool>
+
 <!-- <ifarchitecture name="!_ppc64le_">
 <bin name="testTFAOT" file="testRunner.cpp,testAOT.cc">
   <flags DNN_NAME="testAOT_add" />
diff --git a/PhysicsTools/TensorFlow/test/createconstantgraph.py b/PhysicsTools/TensorFlow/test/createconstantgraph.py
index 7eecc62d9bccc..9cddb40a5fc94 100644
--- a/PhysicsTools/TensorFlow/test/createconstantgraph.py
+++ b/PhysicsTools/TensorFlow/test/createconstantgraph.py
@@ -34,7 +34,12 @@
 h = tf.add(tf.matmul(x_, W), b)
 y = tf.multiply(h, scale_, name="output")
 
-sess = tf.Session()
+
+# Setup the script to run on CPU only
+config = tf.ConfigProto(
+        device_count = {'GPU': 0}
+    )
+sess = tf.Session(config=config)
 sess.run(tf.global_variables_initializer())
 
 print(sess.run(y, feed_dict={scale_: 1.0, x_: [range(10)]})[0][0])
diff --git a/PhysicsTools/TensorFlow/test/creategraph.py b/PhysicsTools/TensorFlow/test/creategraph.py
index b9b5352c2a422..9d6195c8185c3 100644
--- a/PhysicsTools/TensorFlow/test/creategraph.py
+++ b/PhysicsTools/TensorFlow/test/creategraph.py
@@ -34,7 +34,11 @@
 h = tf.add(tf.matmul(x_, W), b)
 y = tf.multiply(h, scale_, name="output")
 
-sess = tf.Session()
+# Setup the script to run on CPU only
+config = tf.ConfigProto(
+        device_count = {'GPU': 0}
+    )
+sess = tf.Session(config=config)
 sess.run(tf.global_variables_initializer())
 
 print(sess.run(y, feed_dict={scale_: 1.0, x_: [range(10)]})[0][0])
diff --git a/PhysicsTools/TensorFlow/test/testBase.h b/PhysicsTools/TensorFlow/test/testBase.h
index 4ef52901c5e2a..914b9b8631ac3 100644
--- a/PhysicsTools/TensorFlow/test/testBase.h
+++ b/PhysicsTools/TensorFlow/test/testBase.h
@@ -20,7 +20,8 @@ class testBase : public CppUnit::TestFixture {
   void tearDown();
   std::string cmsswPath(std::string path);
 
-  virtual void checkAll() = 0;
+  virtual void test() = 0;
+
   virtual std::string pyScript() const = 0;
 };
 
diff --git a/PhysicsTools/TensorFlow/test/testBaseCUDA.h b/PhysicsTools/TensorFlow/test/testBaseCUDA.h
new file mode 100644
index 0000000000000..6d2fda4277ebc
--- /dev/null
+++ b/PhysicsTools/TensorFlow/test/testBaseCUDA.h
@@ -0,0 +1,89 @@
+/*
+ * Base class for tests.
+ *
+ * Author: Marcel Rieger
+ */
+
+#ifndef PHYSICSTOOLS_TENSORFLOW_TEST_TESTBASECUDA_H
+#define PHYSICSTOOLS_TENSORFLOW_TEST_TESTBASECUDA_H
+
+#include <boost/filesystem.hpp>
+#include <filesystem>
+#include <cppunit/extensions/HelperMacros.h>
+#include <stdexcept>
+#include "catch.hpp"
+
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
+#include "FWCore/ParameterSetReader/interface/ParameterSetReader.h"
+#include "FWCore/PluginManager/interface/PluginManager.h"
+#include "FWCore/PluginManager/interface/standard.h"
+#include "FWCore/ServiceRegistry/interface/Service.h"
+#include "FWCore/ServiceRegistry/interface/ServiceRegistry.h"
+#include "FWCore/ServiceRegistry/interface/ServiceToken.h"
+#include "FWCore/Utilities/interface/Exception.h"
+#include "FWCore/Utilities/interface/ResourceInformation.h"
+#include "HeterogeneousCore/CUDAServices/interface/CUDAService.h"
+#include "HeterogeneousCore/CUDAUtilities/interface/requireDevices.h"
+
+namespace {
+  CUDAService makeCUDAService(edm::ParameterSet ps) {
+    auto desc = edm::ConfigurationDescriptions("Service", "CUDAService");
+    CUDAService::fillDescriptions(desc);
+    desc.validate(ps, "CUDAService");
+    return CUDAService(ps);
+  }
+}  // namespace
+
+class testBaseCUDA : public CppUnit::TestFixture {
+public:
+  std::string dataPath_;
+
+  void setUp();
+  void tearDown();
+  std::string cmsswPath(std::string path);
+
+  virtual void test() = 0;
+
+  virtual std::string pyScript() const = 0;
+};
+
+void testBaseCUDA::setUp() {
+  dataPath_ =
+      cmsswPath("/test/" + std::string(std::getenv("SCRAM_ARCH")) + "/" + boost::filesystem::unique_path().string());
+
+  // create the graph
+  std::string testPath = cmsswPath("/src/PhysicsTools/TensorFlow/test");
+  std::string cmd = "python3 -W ignore " + testPath + "/" + pyScript() + " " + dataPath_;
+  std::array<char, 128> buffer;
+  std::string result;
+  std::shared_ptr<FILE> pipe(popen(cmd.c_str(), "r"), pclose);
+  if (!pipe) {
+    throw std::runtime_error("popen() failed!");
+  }
+  while (!feof(pipe.get())) {
+    if (fgets(buffer.data(), 128, pipe.get()) != NULL) {
+      result += buffer.data();
+    }
+  }
+  std::cout << std::endl << result << std::endl;
+}
+
+void testBaseCUDA::tearDown() {
+  if (std::filesystem::exists(dataPath_)) {
+    std::filesystem::remove_all(dataPath_);
+  }
+}
+
+std::string testBaseCUDA::cmsswPath(std::string path) {
+  if (path.size() > 0 && path.substr(0, 1) != "/") {
+    path = "/" + path;
+  }
+
+  std::string base = std::string(std::getenv("CMSSW_BASE"));
+  std::string releaseBase = std::string(std::getenv("CMSSW_RELEASE_BASE"));
+
+  return (std::filesystem::exists(base.c_str()) ? base : releaseBase) + path;
+}
+
+#endif  // PHYSICSTOOLS_TENSORFLOW_TEST_TESTBASECUDA_H
diff --git a/PhysicsTools/TensorFlow/test/testConstSession.cc b/PhysicsTools/TensorFlow/test/testConstSession.cc
index 315c97daf1ddb..7fdf8e36c823b 100644
--- a/PhysicsTools/TensorFlow/test/testConstSession.cc
+++ b/PhysicsTools/TensorFlow/test/testConstSession.cc
@@ -14,28 +14,31 @@
 
 class testConstSession : public testBase {
   CPPUNIT_TEST_SUITE(testConstSession);
-  CPPUNIT_TEST(checkAll);
+  CPPUNIT_TEST(test);
   CPPUNIT_TEST_SUITE_END();
 
 public:
   std::string pyScript() const override;
-  void checkAll() override;
+  void test() override;
 };
 
 CPPUNIT_TEST_SUITE_REGISTRATION(testConstSession);
 
 std::string testConstSession::pyScript() const { return "createconstantgraph.py"; }
 
-void testConstSession::checkAll() {
+void testConstSession::test() {
   std::string pbFile = dataPath_ + "/constantgraph.pb";
 
+  std::cout << "Testing CPU backend" << std::endl;
+  tensorflow::Backend backend = tensorflow::Backend::cpu;
+
   // load the graph
   tensorflow::setLogging();
   tensorflow::GraphDef* graphDef = tensorflow::loadGraphDef(pbFile);
   CPPUNIT_ASSERT(graphDef != nullptr);
 
   // create a new session and add the graphDef
-  const tensorflow::Session* session = tensorflow::createSession(graphDef);
+  const tensorflow::Session* session = tensorflow::createSession(graphDef, backend);
   CPPUNIT_ASSERT(session != nullptr);
 
   // example evaluation
diff --git a/PhysicsTools/TensorFlow/test/testConstSessionCUDA.cc b/PhysicsTools/TensorFlow/test/testConstSessionCUDA.cc
new file mode 100644
index 0000000000000..1236585451deb
--- /dev/null
+++ b/PhysicsTools/TensorFlow/test/testConstSessionCUDA.cc
@@ -0,0 +1,91 @@
+/*
+ * Tests for working with constant sessions.
+ * For more info, see https://gitlab.cern.ch/mrieger/CMSSW-DNN.
+ *
+ * Author: Marcel Rieger
+ */
+
+#include <stdexcept>
+#include <cppunit/extensions/HelperMacros.h>
+
+#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+
+#include "testBaseCUDA.h"
+
+class testConstSessionCUDA : public testBaseCUDA {
+  CPPUNIT_TEST_SUITE(testConstSessionCUDA);
+  CPPUNIT_TEST(test);
+  CPPUNIT_TEST_SUITE_END();
+
+public:
+  std::string pyScript() const override;
+  void test() override;
+};
+
+CPPUNIT_TEST_SUITE_REGISTRATION(testConstSessionCUDA);
+
+std::string testConstSessionCUDA::pyScript() const { return "createconstantgraph.py"; }
+
+void testConstSessionCUDA::test() {
+  if (!cms::cudatest::testDevices())
+    return;
+
+  std::vector<edm::ParameterSet> psets;
+  edm::ServiceToken serviceToken = edm::ServiceRegistry::createSet(psets);
+  edm::ServiceRegistry::Operate operate(serviceToken);
+
+  // Setup the CUDA Service
+  edmplugin::PluginManager::configure(edmplugin::standard::config());
+
+  std::string const config = R"_(import FWCore.ParameterSet.Config as cms
+process = cms.Process('Test')
+process.add_(cms.Service('ResourceInformationService'))
+process.add_(cms.Service('CUDAService'))
+)_";
+  std::unique_ptr<edm::ParameterSet> params;
+  edm::makeParameterSets(config, params);
+  edm::ServiceToken tempToken(edm::ServiceRegistry::createServicesFromConfig(std::move(params)));
+  edm::ServiceRegistry::Operate operate2(tempToken);
+
+  auto cs = makeCUDAService(edm::ParameterSet{});
+  std::cout << "CUDA service enabled: " << cs.enabled() << std::endl;
+
+  std::string pbFile = dataPath_ + "/constantgraph.pb";
+
+  std::cout << "Testing CUDA backend" << std::endl;
+  tensorflow::Backend backend = tensorflow::Backend::cuda;
+
+  // load the graph
+  tensorflow::setLogging();
+  tensorflow::GraphDef* graphDef = tensorflow::loadGraphDef(pbFile);
+  CPPUNIT_ASSERT(graphDef != nullptr);
+
+  // create a new session and add the graphDef
+  const tensorflow::Session* session = tensorflow::createSession(graphDef, backend);
+  CPPUNIT_ASSERT(session != nullptr);
+
+  // example evaluation
+  tensorflow::Tensor input(tensorflow::DT_FLOAT, {1, 10});
+  float* d = input.flat<float>().data();
+  for (size_t i = 0; i < 10; i++, d++) {
+    *d = float(i);
+  }
+  tensorflow::Tensor scale(tensorflow::DT_FLOAT, {});
+  scale.scalar<float>()() = 1.0;
+  std::vector<tensorflow::Tensor> outputs;
+
+  // run using the convenience helper
+  outputs.clear();
+  tensorflow::run(session, {{"input", input}, {"scale", scale}}, {"output"}, &outputs);
+  CPPUNIT_ASSERT(outputs.size() == 1);
+  std::cout << outputs[0].DebugString() << std::endl;
+  CPPUNIT_ASSERT(outputs[0].matrix<float>()(0, 0) == 46.);
+
+  // check for exception
+  CPPUNIT_ASSERT_THROW(tensorflow::run(session, {{"foo", input}}, {"output"}, &outputs), cms::Exception);
+
+  // cleanup
+  CPPUNIT_ASSERT(tensorflow::closeSession(session));
+  CPPUNIT_ASSERT(session == nullptr);
+  delete graphDef;
+}
diff --git a/PhysicsTools/TensorFlow/test/testGraphLoading.cc b/PhysicsTools/TensorFlow/test/testGraphLoading.cc
index 1f01c87d7ec57..95047ec7c04aa 100644
--- a/PhysicsTools/TensorFlow/test/testGraphLoading.cc
+++ b/PhysicsTools/TensorFlow/test/testGraphLoading.cc
@@ -14,32 +14,35 @@
 
 class testGraphLoading : public testBase {
   CPPUNIT_TEST_SUITE(testGraphLoading);
-  CPPUNIT_TEST(checkAll);
+  CPPUNIT_TEST(test);
   CPPUNIT_TEST_SUITE_END();
 
 public:
   std::string pyScript() const override;
-  void checkAll() override;
+  void test() override;
 };
 
 CPPUNIT_TEST_SUITE_REGISTRATION(testGraphLoading);
 
 std::string testGraphLoading::pyScript() const { return "createconstantgraph.py"; }
 
-void testGraphLoading::checkAll() {
+void testGraphLoading::test() {
   std::string pbFile = dataPath_ + "/constantgraph.pb";
 
+  std::cout << "Testing CPU backend" << std::endl;
+  tensorflow::Backend backend = tensorflow::Backend::cpu;
+
   // load the graph
   tensorflow::setLogging();
   tensorflow::GraphDef* graphDef = tensorflow::loadGraphDef(pbFile);
   CPPUNIT_ASSERT(graphDef != nullptr);
 
   // create a new session and add the graphDef
-  tensorflow::Session* session = tensorflow::createSession(graphDef);
+  tensorflow::Session* session = tensorflow::createSession(graphDef, backend);
   CPPUNIT_ASSERT(session != nullptr);
 
   // check for exception
-  CPPUNIT_ASSERT_THROW(tensorflow::createSession(nullptr), cms::Exception);
+  CPPUNIT_ASSERT_THROW(tensorflow::createSession(nullptr, backend), cms::Exception);
 
   // example evaluation
   tensorflow::Tensor input(tensorflow::DT_FLOAT, {1, 10});
diff --git a/PhysicsTools/TensorFlow/test/testGraphLoadingCUDA.cc b/PhysicsTools/TensorFlow/test/testGraphLoadingCUDA.cc
new file mode 100644
index 0000000000000..233ac6b9c2ebb
--- /dev/null
+++ b/PhysicsTools/TensorFlow/test/testGraphLoadingCUDA.cc
@@ -0,0 +1,105 @@
+/*
+ * Tests for loading graphs via the converted protobuf files.
+ * For more info, see https://gitlab.cern.ch/mrieger/CMSSW-DNN.
+ *
+ * Author: Marcel Rieger
+ */
+
+#include <stdexcept>
+#include <cppunit/extensions/HelperMacros.h>
+
+#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+
+#include "testBaseCUDA.h"
+
+class testGraphLoadingCUDA : public testBaseCUDA {
+  CPPUNIT_TEST_SUITE(testGraphLoadingCUDA);
+  CPPUNIT_TEST(test);
+  CPPUNIT_TEST_SUITE_END();
+
+public:
+  std::string pyScript() const override;
+  void test() override;
+};
+
+CPPUNIT_TEST_SUITE_REGISTRATION(testGraphLoadingCUDA);
+
+std::string testGraphLoadingCUDA::pyScript() const { return "createconstantgraph.py"; }
+
+void testGraphLoadingCUDA::test() {
+  if (!cms::cudatest::testDevices())
+    return;
+
+  std::vector<edm::ParameterSet> psets;
+  edm::ServiceToken serviceToken = edm::ServiceRegistry::createSet(psets);
+  edm::ServiceRegistry::Operate operate(serviceToken);
+
+  // Setup the CUDA Service
+  edmplugin::PluginManager::configure(edmplugin::standard::config());
+
+  std::string const config = R"_(import FWCore.ParameterSet.Config as cms
+process = cms.Process('Test')
+process.add_(cms.Service('ResourceInformationService'))
+process.add_(cms.Service('CUDAService'))
+)_";
+  std::unique_ptr<edm::ParameterSet> params;
+  edm::makeParameterSets(config, params);
+  edm::ServiceToken tempToken(edm::ServiceRegistry::createServicesFromConfig(std::move(params)));
+  edm::ServiceRegistry::Operate operate2(tempToken);
+
+  auto cs = makeCUDAService(edm::ParameterSet{});
+  std::cout << "CUDA service enabled: " << cs.enabled() << std::endl;
+
+  std::string pbFile = dataPath_ + "/constantgraph.pb";
+
+  std::cout << "Testing CUDA backend" << std::endl;
+  tensorflow::Backend backend = tensorflow::Backend::cuda;
+
+  // load the graph
+  tensorflow::setLogging();
+  tensorflow::GraphDef* graphDef = tensorflow::loadGraphDef(pbFile);
+  CPPUNIT_ASSERT(graphDef != nullptr);
+
+  // create a new session and add the graphDef
+  tensorflow::Session* session = tensorflow::createSession(graphDef, backend);
+  CPPUNIT_ASSERT(session != nullptr);
+
+  // check for exception
+  CPPUNIT_ASSERT_THROW(tensorflow::createSession(nullptr, backend), cms::Exception);
+
+  // example evaluation
+  tensorflow::Tensor input(tensorflow::DT_FLOAT, {1, 10});
+  float* d = input.flat<float>().data();
+  for (size_t i = 0; i < 10; i++, d++) {
+    *d = float(i);
+  }
+  tensorflow::Tensor scale(tensorflow::DT_FLOAT, {});
+  scale.scalar<float>()() = 1.0;
+
+  std::vector<tensorflow::Tensor> outputs;
+  tensorflow::Status status = session->Run({{"input", input}, {"scale", scale}}, {"output"}, {}, &outputs);
+  if (!status.ok()) {
+    std::cout << status.ToString() << std::endl;
+    CPPUNIT_ASSERT(false);
+  }
+
+  // check the output
+  CPPUNIT_ASSERT(outputs.size() == 1);
+  std::cout << outputs[0].DebugString() << std::endl;
+  CPPUNIT_ASSERT(outputs[0].matrix<float>()(0, 0) == 46.);
+
+  // run again using the convenience helper
+  outputs.clear();
+  tensorflow::run(session, {{"input", input}, {"scale", scale}}, {"output"}, &outputs);
+  CPPUNIT_ASSERT(outputs.size() == 1);
+  std::cout << outputs[0].DebugString() << std::endl;
+  CPPUNIT_ASSERT(outputs[0].matrix<float>()(0, 0) == 46.);
+
+  // check for exception
+  CPPUNIT_ASSERT_THROW(tensorflow::run(session, {{"foo", input}}, {"output"}, &outputs), cms::Exception);
+
+  // cleanup
+  CPPUNIT_ASSERT(tensorflow::closeSession(session));
+  CPPUNIT_ASSERT(session == nullptr);
+  delete graphDef;
+}
diff --git a/PhysicsTools/TensorFlow/test/testHelloWorld.cc b/PhysicsTools/TensorFlow/test/testHelloWorld.cc
index cd6da31c0dd43..bf1161b3d5c60 100644
--- a/PhysicsTools/TensorFlow/test/testHelloWorld.cc
+++ b/PhysicsTools/TensorFlow/test/testHelloWorld.cc
@@ -10,29 +10,34 @@
 
 #include "tensorflow/cc/saved_model/loader.h"
 #include "tensorflow/cc/saved_model/tag_constants.h"
+#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
 
 #include "testBase.h"
 
 class testHelloWorld : public testBase {
   CPPUNIT_TEST_SUITE(testHelloWorld);
-  CPPUNIT_TEST(checkAll);
+  CPPUNIT_TEST(test);
   CPPUNIT_TEST_SUITE_END();
 
 public:
   std::string pyScript() const override;
-  void checkAll() override;
+  void test() override;
 };
 
 CPPUNIT_TEST_SUITE_REGISTRATION(testHelloWorld);
 
 std::string testHelloWorld::pyScript() const { return "creategraph.py"; }
 
-void testHelloWorld::checkAll() {
+void testHelloWorld::test() {
   std::string modelDir = dataPath_ + "/simplegraph";
+  // Testing CPU
+  std::cout << "Testing CPU backend" << std::endl;
+  tensorflow::Backend backend = tensorflow::Backend::cpu;
 
   // object to load and run the graph / session
   tensorflow::Status status;
   tensorflow::SessionOptions sessionOptions;
+  tensorflow::setBackend(sessionOptions, backend);
   tensorflow::RunOptions runOptions;
   tensorflow::SavedModelBundle bundle;
 
diff --git a/PhysicsTools/TensorFlow/test/testHelloWorldCUDA.cc b/PhysicsTools/TensorFlow/test/testHelloWorldCUDA.cc
new file mode 100644
index 0000000000000..0f7070f3a72c7
--- /dev/null
+++ b/PhysicsTools/TensorFlow/test/testHelloWorldCUDA.cc
@@ -0,0 +1,105 @@
+/*
+ * HelloWorld test of the TensorFlow interface.
+ * For more info, see https://gitlab.cern.ch/mrieger/CMSSW-DNN.
+ *
+ * Author: Marcel Rieger
+ */
+
+#include <stdexcept>
+#include <cppunit/extensions/HelperMacros.h>
+
+#include "tensorflow/cc/saved_model/loader.h"
+#include "tensorflow/cc/saved_model/tag_constants.h"
+#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+
+#include "testBaseCUDA.h"
+
+class testHelloWorldCUDA : public testBaseCUDA {
+  CPPUNIT_TEST_SUITE(testHelloWorldCUDA);
+  CPPUNIT_TEST(test);
+  CPPUNIT_TEST_SUITE_END();
+
+public:
+  std::string pyScript() const override;
+  void test() override;
+};
+
+CPPUNIT_TEST_SUITE_REGISTRATION(testHelloWorldCUDA);
+
+std::string testHelloWorldCUDA::pyScript() const { return "creategraph.py"; }
+
+void testHelloWorldCUDA::test() {
+  if (!cms::cudatest::testDevices())
+    return;
+
+  std::vector<edm::ParameterSet> psets;
+  edm::ServiceToken serviceToken = edm::ServiceRegistry::createSet(psets);
+  edm::ServiceRegistry::Operate operate(serviceToken);
+
+  // Setup the CUDA Service
+  edmplugin::PluginManager::configure(edmplugin::standard::config());
+
+  std::string const config = R"_(import FWCore.ParameterSet.Config as cms
+process = cms.Process('Test')
+process.add_(cms.Service('ResourceInformationService'))
+process.add_(cms.Service('CUDAService'))
+)_";
+  std::unique_ptr<edm::ParameterSet> params;
+  edm::makeParameterSets(config, params);
+  edm::ServiceToken tempToken(edm::ServiceRegistry::createServicesFromConfig(std::move(params)));
+  edm::ServiceRegistry::Operate operate2(tempToken);
+
+  auto cs = makeCUDAService(edm::ParameterSet{});
+  std::cout << "CUDA service enabled: " << cs.enabled() << std::endl;
+
+  std::string modelDir = dataPath_ + "/simplegraph";
+  // Testing CPU
+  std::cout << "Testing CUDA backend" << std::endl;
+  tensorflow::Backend backend = tensorflow::Backend::cuda;
+
+  // object to load and run the graph / session
+  tensorflow::Status status;
+  tensorflow::SessionOptions sessionOptions;
+  tensorflow::setBackend(sessionOptions, backend);
+  tensorflow::RunOptions runOptions;
+  tensorflow::SavedModelBundle bundle;
+
+  // load everything
+  status = tensorflow::LoadSavedModel(sessionOptions, runOptions, modelDir, {"serve"}, &bundle);
+  if (!status.ok()) {
+    std::cout << status.ToString() << std::endl;
+    return;
+  }
+
+  // fetch the session
+  tensorflow::Session* session = bundle.session.release();
+
+  // prepare inputs
+  tensorflow::Tensor input(tensorflow::DT_FLOAT, {1, 10});
+  float* d = input.flat<float>().data();
+  for (size_t i = 0; i < 10; i++, d++) {
+    *d = float(i);
+  }
+  tensorflow::Tensor scale(tensorflow::DT_FLOAT, {});
+  scale.scalar<float>()() = 1.0;
+
+  // prepare outputs
+  std::vector<tensorflow::Tensor> outputs;
+
+  // session run
+  status = session->Run({{"input", input}, {"scale", scale}}, {"output"}, {}, &outputs);
+  if (!status.ok()) {
+    std::cout << status.ToString() << std::endl;
+    return;
+  }
+
+  // log the output tensor
+  std::cout << outputs[0].DebugString() << std::endl;
+
+  // close the session
+  status = session->Close();
+  if (!status.ok()) {
+    std::cerr << "error while closing session" << std::endl;
+  }
+  delete session;
+}
diff --git a/PhysicsTools/TensorFlow/test/testMetaGraphLoading.cc b/PhysicsTools/TensorFlow/test/testMetaGraphLoading.cc
index 4868207c5616a..ac9b7406b4a43 100644
--- a/PhysicsTools/TensorFlow/test/testMetaGraphLoading.cc
+++ b/PhysicsTools/TensorFlow/test/testMetaGraphLoading.cc
@@ -14,36 +14,39 @@
 
 class testMetaGraphLoading : public testBase {
   CPPUNIT_TEST_SUITE(testMetaGraphLoading);
-  CPPUNIT_TEST(checkAll);
+  CPPUNIT_TEST(test);
   CPPUNIT_TEST_SUITE_END();
 
 public:
   std::string pyScript() const override;
-  void checkAll() override;
+  void test() override;
 };
 
 CPPUNIT_TEST_SUITE_REGISTRATION(testMetaGraphLoading);
 
 std::string testMetaGraphLoading::pyScript() const { return "creategraph.py"; }
 
-void testMetaGraphLoading::checkAll() {
+void testMetaGraphLoading::test() {
   std::string exportDir = dataPath_ + "/simplegraph";
 
+  std::cout << "Testing CPU backend" << std::endl;
+  tensorflow::Backend backend = tensorflow::Backend::cpu;
+
   // load the graph
   tensorflow::setLogging();
   tensorflow::MetaGraphDef* metaGraphDef = tensorflow::loadMetaGraphDef(exportDir);
   CPPUNIT_ASSERT(metaGraphDef != nullptr);
 
   // create a new, empty session
-  tensorflow::Session* session1 = tensorflow::createSession();
+  tensorflow::Session* session1 = tensorflow::createSession(backend);
   CPPUNIT_ASSERT(session1 != nullptr);
 
   // create a new session, using the meta graph
-  tensorflow::Session* session2 = tensorflow::createSession(metaGraphDef, exportDir);
+  tensorflow::Session* session2 = tensorflow::createSession(metaGraphDef, exportDir, backend);
   CPPUNIT_ASSERT(session2 != nullptr);
 
   // check for exception
-  CPPUNIT_ASSERT_THROW(tensorflow::createSession(nullptr, exportDir), cms::Exception);
+  CPPUNIT_ASSERT_THROW(tensorflow::createSession(nullptr, exportDir, backend), cms::Exception);
 
   // example evaluation
   tensorflow::Tensor input(tensorflow::DT_FLOAT, {1, 10});
diff --git a/PhysicsTools/TensorFlow/test/testMetaGraphLoadingCUDA.cc b/PhysicsTools/TensorFlow/test/testMetaGraphLoadingCUDA.cc
new file mode 100644
index 0000000000000..11ac4dd1405f5
--- /dev/null
+++ b/PhysicsTools/TensorFlow/test/testMetaGraphLoadingCUDA.cc
@@ -0,0 +1,108 @@
+/*
+ * Tests for loading meta graphs via the SavedModel interface.
+ * For more info, see https://gitlab.cern.ch/mrieger/CMSSW-DNN.
+ *
+ * Author: Marcel Rieger
+ */
+
+#include <stdexcept>
+#include <cppunit/extensions/HelperMacros.h>
+
+#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+
+#include "testBaseCUDA.h"
+
+class testMetaGraphLoadingCUDA : public testBaseCUDA {
+  CPPUNIT_TEST_SUITE(testMetaGraphLoadingCUDA);
+  CPPUNIT_TEST(test);
+  CPPUNIT_TEST_SUITE_END();
+
+public:
+  std::string pyScript() const override;
+  void test() override;
+};
+
+CPPUNIT_TEST_SUITE_REGISTRATION(testMetaGraphLoadingCUDA);
+
+std::string testMetaGraphLoadingCUDA::pyScript() const { return "creategraph.py"; }
+
+void testMetaGraphLoadingCUDA::test() {
+  if (!cms::cudatest::testDevices())
+    return;
+
+  std::vector<edm::ParameterSet> psets;
+  edm::ServiceToken serviceToken = edm::ServiceRegistry::createSet(psets);
+  edm::ServiceRegistry::Operate operate(serviceToken);
+
+  // Setup the CUDA Service
+  edmplugin::PluginManager::configure(edmplugin::standard::config());
+
+  std::string const config = R"_(import FWCore.ParameterSet.Config as cms
+process = cms.Process('Test')
+process.add_(cms.Service('ResourceInformationService'))
+process.add_(cms.Service('CUDAService'))
+)_";
+  std::unique_ptr<edm::ParameterSet> params;
+  edm::makeParameterSets(config, params);
+  edm::ServiceToken tempToken(edm::ServiceRegistry::createServicesFromConfig(std::move(params)));
+  edm::ServiceRegistry::Operate operate2(tempToken);
+
+  auto cs = makeCUDAService(edm::ParameterSet{});
+  std::cout << "CUDA service enabled: " << cs.enabled() << std::endl;
+  std::string exportDir = dataPath_ + "/simplegraph";
+
+  std::cout << "Testing CUDA backend" << std::endl;
+  tensorflow::Backend backend = tensorflow::Backend::cuda;
+
+  // load the graph
+  tensorflow::setLogging();
+  tensorflow::MetaGraphDef* metaGraphDef = tensorflow::loadMetaGraphDef(exportDir);
+  CPPUNIT_ASSERT(metaGraphDef != nullptr);
+
+  // create a new, empty session
+  tensorflow::Session* session1 = tensorflow::createSession(backend);
+  CPPUNIT_ASSERT(session1 != nullptr);
+
+  // create a new session, using the meta graph
+  tensorflow::Session* session2 = tensorflow::createSession(metaGraphDef, exportDir, backend);
+  CPPUNIT_ASSERT(session2 != nullptr);
+
+  // check for exception
+  CPPUNIT_ASSERT_THROW(tensorflow::createSession(nullptr, exportDir, backend), cms::Exception);
+
+  // example evaluation
+  tensorflow::Tensor input(tensorflow::DT_FLOAT, {1, 10});
+  float* d = input.flat<float>().data();
+  for (size_t i = 0; i < 10; i++, d++) {
+    *d = float(i);
+  }
+  tensorflow::Tensor scale(tensorflow::DT_FLOAT, {});
+  scale.scalar<float>()() = 1.0;
+
+  std::vector<tensorflow::Tensor> outputs;
+  tensorflow::Status status = session2->Run({{"input", input}, {"scale", scale}}, {"output"}, {}, &outputs);
+  if (!status.ok()) {
+    std::cout << status.ToString() << std::endl;
+    CPPUNIT_ASSERT(false);
+  }
+
+  // check the output
+  CPPUNIT_ASSERT(outputs.size() == 1);
+  std::cout << outputs[0].DebugString() << std::endl;
+  CPPUNIT_ASSERT(outputs[0].matrix<float>()(0, 0) == 46.);
+
+  // run again using the convenience helper
+  outputs.clear();
+  tensorflow::run(session2, {{"input", input}, {"scale", scale}}, {"output"}, &outputs);
+  CPPUNIT_ASSERT(outputs.size() == 1);
+  std::cout << outputs[0].DebugString() << std::endl;
+  CPPUNIT_ASSERT(outputs[0].matrix<float>()(0, 0) == 46.);
+
+  // check for exception
+  CPPUNIT_ASSERT_THROW(tensorflow::run(session2, {{"foo", input}}, {"output"}, &outputs), cms::Exception);
+
+  // cleanup
+  CPPUNIT_ASSERT(tensorflow::closeSession(session1));
+  CPPUNIT_ASSERT(tensorflow::closeSession(session2));
+  delete metaGraphDef;
+}
diff --git a/PhysicsTools/TensorFlow/test/testSessionCache.cc b/PhysicsTools/TensorFlow/test/testSessionCache.cc
index ba241402744f6..9cec510b4cea8 100644
--- a/PhysicsTools/TensorFlow/test/testSessionCache.cc
+++ b/PhysicsTools/TensorFlow/test/testSessionCache.cc
@@ -13,25 +13,28 @@
 
 class testSessionCache : public testBase {
   CPPUNIT_TEST_SUITE(testSessionCache);
-  CPPUNIT_TEST(checkAll);
+  CPPUNIT_TEST(test);
   CPPUNIT_TEST_SUITE_END();
 
 public:
   std::string pyScript() const override;
-  void checkAll() override;
+  void test() override;
 };
 
 CPPUNIT_TEST_SUITE_REGISTRATION(testSessionCache);
 
 std::string testSessionCache::pyScript() const { return "createconstantgraph.py"; }
 
-void testSessionCache::checkAll() {
+void testSessionCache::test() {
   std::string pbFile = dataPath_ + "/constantgraph.pb";
 
+  std::cout << "Testing CPU backend" << std::endl;
+  tensorflow::Backend backend = tensorflow::Backend::cpu;
+
   tensorflow::setLogging();
 
   // load the graph and the session
-  tensorflow::SessionCache cache(pbFile);
+  tensorflow::SessionCache cache(pbFile, backend);
   CPPUNIT_ASSERT(cache.graph.load() != nullptr);
   CPPUNIT_ASSERT(cache.session.load() != nullptr);
 
diff --git a/PhysicsTools/TensorFlow/test/testSessionCacheCUDA.cc b/PhysicsTools/TensorFlow/test/testSessionCacheCUDA.cc
new file mode 100644
index 0000000000000..fcb0952ce643a
--- /dev/null
+++ b/PhysicsTools/TensorFlow/test/testSessionCacheCUDA.cc
@@ -0,0 +1,72 @@
+/*
+ * Tests for interacting with the SessionCache.
+ *
+ * Author: Marcel Rieger
+ */
+
+#include <stdexcept>
+#include <cppunit/extensions/HelperMacros.h>
+
+#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+
+#include "testBaseCUDA.h"
+
+class testSessionCacheCUDA : public testBaseCUDA {
+  CPPUNIT_TEST_SUITE(testSessionCacheCUDA);
+  CPPUNIT_TEST(test);
+  CPPUNIT_TEST_SUITE_END();
+
+public:
+  std::string pyScript() const override;
+  void test() override;
+};
+
+CPPUNIT_TEST_SUITE_REGISTRATION(testSessionCacheCUDA);
+
+std::string testSessionCacheCUDA::pyScript() const { return "createconstantgraph.py"; }
+
+void testSessionCacheCUDA::test() {
+  if (!cms::cudatest::testDevices())
+    return;
+
+  std::vector<edm::ParameterSet> psets;
+  edm::ServiceToken serviceToken = edm::ServiceRegistry::createSet(psets);
+  edm::ServiceRegistry::Operate operate(serviceToken);
+
+  // Setup the CUDA Service
+  edmplugin::PluginManager::configure(edmplugin::standard::config());
+
+  std::string const config = R"_(import FWCore.ParameterSet.Config as cms
+process = cms.Process('Test')
+process.add_(cms.Service('ResourceInformationService'))
+process.add_(cms.Service('CUDAService'))
+)_";
+  std::unique_ptr<edm::ParameterSet> params;
+  edm::makeParameterSets(config, params);
+  edm::ServiceToken tempToken(edm::ServiceRegistry::createServicesFromConfig(std::move(params)));
+  edm::ServiceRegistry::Operate operate2(tempToken);
+
+  auto cs = makeCUDAService(edm::ParameterSet{});
+  std::cout << "CUDA service enabled: " << cs.enabled() << std::endl;
+
+  std::string pbFile = dataPath_ + "/constantgraph.pb";
+
+  std::cout << "Testing CUDA backend" << std::endl;
+  tensorflow::Backend backend = tensorflow::Backend::cuda;
+
+  tensorflow::setLogging();
+
+  // load the graph and the session
+  tensorflow::SessionCache cache(pbFile, backend);
+  CPPUNIT_ASSERT(cache.graph.load() != nullptr);
+  CPPUNIT_ASSERT(cache.session.load() != nullptr);
+
+  // get a const session pointer
+  const tensorflow::Session* session = cache.getSession();
+  CPPUNIT_ASSERT(session != nullptr);
+
+  // cleanup
+  cache.closeSession();
+  CPPUNIT_ASSERT(cache.graph.load() == nullptr);
+  CPPUNIT_ASSERT(cache.session.load() == nullptr);
+}
diff --git a/PhysicsTools/TensorFlow/test/testThreadPools.cc b/PhysicsTools/TensorFlow/test/testThreadPools.cc
index f9f03764e1ada..6c07b1d3c78ef 100644
--- a/PhysicsTools/TensorFlow/test/testThreadPools.cc
+++ b/PhysicsTools/TensorFlow/test/testThreadPools.cc
@@ -14,21 +14,24 @@
 
 class testGraphLoading : public testBase {
   CPPUNIT_TEST_SUITE(testGraphLoading);
-  CPPUNIT_TEST(checkAll);
+  CPPUNIT_TEST(test);
   CPPUNIT_TEST_SUITE_END();
 
 public:
   std::string pyScript() const override;
-  void checkAll() override;
+  void test() override;
 };
 
 CPPUNIT_TEST_SUITE_REGISTRATION(testGraphLoading);
 
 std::string testGraphLoading::pyScript() const { return "createconstantgraph.py"; }
 
-void testGraphLoading::checkAll() {
+void testGraphLoading::test() {
   std::string pbFile = dataPath_ + "/constantgraph.pb";
 
+  std::cout << "Testing CPU backend" << std::endl;
+  tensorflow::Backend backend = tensorflow::Backend::cpu;
+
   // initialize the TBB threadpool
   int nThreads = 4;
   tensorflow::TBBThreadPool::instance(nThreads);
@@ -39,7 +42,7 @@ void testGraphLoading::checkAll() {
   CPPUNIT_ASSERT(graphDef != nullptr);
 
   // create a new session and add the graphDef
-  tensorflow::Session* session = tensorflow::createSession(graphDef);
+  tensorflow::Session* session = tensorflow::createSession(graphDef, backend);
   CPPUNIT_ASSERT(session != nullptr);
 
   // prepare inputs
@@ -66,7 +69,7 @@ void testGraphLoading::checkAll() {
   CPPUNIT_ASSERT(outputs[0].matrix<float>()(0, 0) == 46.);
 
   // tensorflow defaut pool using a new session
-  tensorflow::Session* session2 = tensorflow::createSession(graphDef, nThreads);
+  tensorflow::Session* session2 = tensorflow::createSession(graphDef, backend, nThreads);
   CPPUNIT_ASSERT(session != nullptr);
   outputs.clear();
   tensorflow::run(session2, {{"input", input}, {"scale", scale}}, {"output"}, &outputs, "tensorflow");
diff --git a/PhysicsTools/TensorFlow/test/testThreadPoolsCUDA.cc b/PhysicsTools/TensorFlow/test/testThreadPoolsCUDA.cc
new file mode 100644
index 0000000000000..333144c19fccb
--- /dev/null
+++ b/PhysicsTools/TensorFlow/test/testThreadPoolsCUDA.cc
@@ -0,0 +1,112 @@
+/*
+ * Tests for running inference using custom thread pools.
+ * For more info, see https://gitlab.cern.ch/mrieger/CMSSW-DNN.
+ *
+ * Author: Marcel Rieger
+ */
+
+#include <cppunit/extensions/HelperMacros.h>
+#include <stdexcept>
+
+#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+
+#include "testBaseCUDA.h"
+
+class testGraphLoadingCUDA : public testBaseCUDA {
+  CPPUNIT_TEST_SUITE(testGraphLoadingCUDA);
+  CPPUNIT_TEST(test);
+  CPPUNIT_TEST_SUITE_END();
+
+public:
+  std::string pyScript() const override;
+  void test() override;
+};
+
+CPPUNIT_TEST_SUITE_REGISTRATION(testGraphLoadingCUDA);
+
+std::string testGraphLoadingCUDA::pyScript() const { return "createconstantgraph.py"; }
+
+void testGraphLoadingCUDA::test() {
+  if (!cms::cudatest::testDevices())
+    return;
+
+  std::vector<edm::ParameterSet> psets;
+  edm::ServiceToken serviceToken = edm::ServiceRegistry::createSet(psets);
+  edm::ServiceRegistry::Operate operate(serviceToken);
+
+  // Setup the CUDA Service
+  edmplugin::PluginManager::configure(edmplugin::standard::config());
+
+  std::string const config = R"_(import FWCore.ParameterSet.Config as cms
+process = cms.Process('Test')
+process.add_(cms.Service('ResourceInformationService'))
+process.add_(cms.Service('CUDAService'))
+)_";
+  std::unique_ptr<edm::ParameterSet> params;
+  edm::makeParameterSets(config, params);
+  edm::ServiceToken tempToken(edm::ServiceRegistry::createServicesFromConfig(std::move(params)));
+  edm::ServiceRegistry::Operate operate2(tempToken);
+
+  auto cs = makeCUDAService(edm::ParameterSet{});
+  std::cout << "CUDA service enabled: " << cs.enabled() << std::endl;
+
+  std::string pbFile = dataPath_ + "/constantgraph.pb";
+
+  std::cout << "Testing CUDA backend" << std::endl;
+  tensorflow::Backend backend = tensorflow::Backend::cuda;
+
+  // initialize the TBB threadpool
+  int nThreads = 4;
+  tensorflow::TBBThreadPool::instance(nThreads);
+
+  // load the graph
+  tensorflow::setLogging();
+  tensorflow::GraphDef* graphDef = tensorflow::loadGraphDef(pbFile);
+  CPPUNIT_ASSERT(graphDef != nullptr);
+
+  // create a new session and add the graphDef
+  tensorflow::Session* session = tensorflow::createSession(graphDef, backend);
+  CPPUNIT_ASSERT(session != nullptr);
+
+  // prepare inputs
+  tensorflow::Tensor input(tensorflow::DT_FLOAT, {1, 10});
+  float* d = input.flat<float>().data();
+  for (size_t i = 0; i < 10; i++, d++) {
+    *d = float(i);
+  }
+  tensorflow::Tensor scale(tensorflow::DT_FLOAT, {});
+  scale.scalar<float>()() = 1.0;
+
+  // "no_threads" pool
+  std::vector<tensorflow::Tensor> outputs;
+  tensorflow::run(session, {{"input", input}, {"scale", scale}}, {"output"}, &outputs, "no_threads");
+  CPPUNIT_ASSERT(outputs.size() == 1);
+  std::cout << outputs[0].DebugString() << std::endl;
+  CPPUNIT_ASSERT(outputs[0].matrix<float>()(0, 0) == 46.);
+
+  // "tbb" pool
+  outputs.clear();
+  tensorflow::run(session, {{"input", input}, {"scale", scale}}, {"output"}, &outputs, "tbb");
+  CPPUNIT_ASSERT(outputs.size() == 1);
+  std::cout << outputs[0].DebugString() << std::endl;
+  CPPUNIT_ASSERT(outputs[0].matrix<float>()(0, 0) == 46.);
+
+  // tensorflow defaut pool using a new session
+  tensorflow::Session* session2 = tensorflow::createSession(graphDef, backend, nThreads);
+  CPPUNIT_ASSERT(session != nullptr);
+  outputs.clear();
+  tensorflow::run(session2, {{"input", input}, {"scale", scale}}, {"output"}, &outputs, "tensorflow");
+  CPPUNIT_ASSERT(outputs.size() == 1);
+  std::cout << outputs[0].DebugString() << std::endl;
+  CPPUNIT_ASSERT(outputs[0].matrix<float>()(0, 0) == 46.);
+
+  // force an exception
+  CPPUNIT_ASSERT_THROW(
+      tensorflow::run(session, {{"input", input}, {"scale", scale}}, {"output"}, &outputs, "not_existing"),
+      cms::Exception);
+
+  // cleanup
+  CPPUNIT_ASSERT(tensorflow::closeSession(session));
+  CPPUNIT_ASSERT(tensorflow::closeSession(session2));
+  delete graphDef;
+}
diff --git a/PhysicsTools/TensorFlow/test/testVisibleDevices.cc b/PhysicsTools/TensorFlow/test/testVisibleDevices.cc
new file mode 100644
index 0000000000000..e1c6a33191925
--- /dev/null
+++ b/PhysicsTools/TensorFlow/test/testVisibleDevices.cc
@@ -0,0 +1,65 @@
+/*
+ * Tests visible device interface 
+ * For more info,
+ * https://github.com/tensorflow/tensorflow/blob/3bc73f5e2ac437b1d9d559751af789c8c965a7f9/tensorflow/core/framework/device_attributes.proto
+ * https://stackoverflow.com/questions/74110853/how-to-check-if-tensorflow-is-using-the-cpu-with-the-c-api\
+ *
+ * Author: Davide Valsecchi
+ */
+
+#include <stdexcept>
+#include <cppunit/extensions/HelperMacros.h>
+
+#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+#include "tensorflow/core/framework/device_attributes.pb.h"
+
+#include "testBase.h"
+
+class testVisibleDevices : public testBase {
+  CPPUNIT_TEST_SUITE(testVisibleDevices);
+  CPPUNIT_TEST(test);
+  CPPUNIT_TEST_SUITE_END();
+
+public:
+  std::string pyScript() const override;
+  void test() override;
+};
+
+CPPUNIT_TEST_SUITE_REGISTRATION(testVisibleDevices);
+
+std::string testVisibleDevices::pyScript() const { return "createconstantgraph.py"; }
+
+void testVisibleDevices::test() {
+  std::string pbFile = dataPath_ + "/constantgraph.pb";
+  tensorflow::Backend backend = tensorflow::Backend::cpu;
+
+  // load the graph
+  tensorflow::setLogging();
+  tensorflow::GraphDef* graphDef = tensorflow::loadGraphDef(pbFile);
+  CPPUNIT_ASSERT(graphDef != nullptr);
+
+  // create a new session and add the graphDef
+  tensorflow::Session* session = tensorflow::createSession(graphDef, backend);
+  CPPUNIT_ASSERT(session != nullptr);
+
+  // check for exception
+  CPPUNIT_ASSERT_THROW(tensorflow::createSession(nullptr, backend), cms::Exception);
+
+  std::vector<tensorflow::DeviceAttributes> response;
+  tensorflow::Status status = session->ListDevices(&response);
+  CPPUNIT_ASSERT(status.ok());
+
+  // If a single device is found, we assume that it's the CPU.
+  // You can check that name if you want to make sure that this is the case
+  std::cout << "Available devices: " << response.size() << std::endl;
+  CPPUNIT_ASSERT(response.size() == 1);
+
+  for (unsigned int i = 0; i < response.size(); ++i) {
+    std::cout << i << " " << response[i].name() << " type: " << response[i].device_type() << std::endl;
+  }
+
+  // cleanup
+  CPPUNIT_ASSERT(tensorflow::closeSession(session));
+  CPPUNIT_ASSERT(session == nullptr);
+  delete graphDef;
+}
diff --git a/PhysicsTools/TensorFlow/test/testVisibleDevicesCUDA.cc b/PhysicsTools/TensorFlow/test/testVisibleDevicesCUDA.cc
new file mode 100644
index 0000000000000..20233dff4fce3
--- /dev/null
+++ b/PhysicsTools/TensorFlow/test/testVisibleDevicesCUDA.cc
@@ -0,0 +1,89 @@
+/*
+ * Tests visible device interface 
+ * For more info,
+ * https://github.com/tensorflow/tensorflow/blob/3bc73f5e2ac437b1d9d559751af789c8c965a7f9/tensorflow/core/framework/device_attributes.proto
+ * https://stackoverflow.com/questions/74110853/how-to-check-if-tensorflow-is-using-the-cpu-with-the-c-api\
+ *
+ * Author: Davide Valsecchi
+ */
+
+#include <stdexcept>
+#include <cppunit/extensions/HelperMacros.h>
+
+#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+#include "tensorflow/core/framework/device_attributes.pb.h"
+
+#include "testBaseCUDA.h"
+
+class testVisibleDevicesCUDA : public testBaseCUDA {
+  CPPUNIT_TEST_SUITE(testVisibleDevicesCUDA);
+  CPPUNIT_TEST(test);
+  CPPUNIT_TEST_SUITE_END();
+
+public:
+  std::string pyScript() const override;
+  void test() override;
+};
+
+CPPUNIT_TEST_SUITE_REGISTRATION(testVisibleDevicesCUDA);
+
+std::string testVisibleDevicesCUDA::pyScript() const { return "createconstantgraph.py"; }
+
+void testVisibleDevicesCUDA::test() {
+  if (!cms::cudatest::testDevices())
+    return;
+
+  std::vector<edm::ParameterSet> psets;
+  edm::ServiceToken serviceToken = edm::ServiceRegistry::createSet(psets);
+  edm::ServiceRegistry::Operate operate(serviceToken);
+
+  // Setup the CUDA Service
+  edmplugin::PluginManager::configure(edmplugin::standard::config());
+
+  std::string const config = R"_(import FWCore.ParameterSet.Config as cms
+process = cms.Process('Test')
+process.add_(cms.Service('ResourceInformationService'))
+process.add_(cms.Service('CUDAService'))
+)_";
+  std::unique_ptr<edm::ParameterSet> params;
+  edm::makeParameterSets(config, params);
+  edm::ServiceToken tempToken(edm::ServiceRegistry::createServicesFromConfig(std::move(params)));
+  edm::ServiceRegistry::Operate operate2(tempToken);
+
+  auto cs = makeCUDAService(edm::ParameterSet{});
+  std::cout << "CUDA service enabled: " << cs.enabled() << std::endl;
+
+  std::string pbFile = dataPath_ + "/constantgraph.pb";
+
+  std::cout << "Testing CUDA backend" << std::endl;
+  tensorflow::Backend backend = tensorflow::Backend::cuda;
+
+  // load the graph
+  tensorflow::setLogging();
+  tensorflow::GraphDef* graphDef = tensorflow::loadGraphDef(pbFile);
+  CPPUNIT_ASSERT(graphDef != nullptr);
+
+  // create a new session and add the graphDef
+  tensorflow::Session* session = tensorflow::createSession(graphDef, backend);
+  CPPUNIT_ASSERT(session != nullptr);
+
+  // check for exception
+  CPPUNIT_ASSERT_THROW(tensorflow::createSession(nullptr, backend), cms::Exception);
+
+  std::vector<tensorflow::DeviceAttributes> response;
+  tensorflow::Status status = session->ListDevices(&response);
+  CPPUNIT_ASSERT(status.ok());
+
+  // If a single device is found, we assume that it's the CPU.
+  // You can check that name if you want to make sure that this is the case
+  std::cout << "Available devices: " << response.size() << std::endl;
+  CPPUNIT_ASSERT(response.size() == 2);
+  for (unsigned int i = 0; i < response.size(); ++i) {
+    std::cout << i << " " << response[i].name() << " type: " << response[i].device_type() << std::endl;
+  }
+
+  // cleanup
+  CPPUNIT_ASSERT(tensorflow::closeSession(session));
+  CPPUNIT_ASSERT(session == nullptr);
+  delete graphDef;
+}
diff --git a/RecoTauTag/RecoTau/src/DeepTauBase.cc b/RecoTauTag/RecoTau/src/DeepTauBase.cc
index 6969db75816fc..2af38735a9c30 100644
--- a/RecoTauTag/RecoTau/src/DeepTauBase.cc
+++ b/RecoTauTag/RecoTau/src/DeepTauBase.cc
@@ -216,6 +216,8 @@ namespace deep_tau {
     for (const auto& graph_entry : graph_names) {
       tensorflow::SessionOptions options;
       tensorflow::setThreading(options, 1);
+      // To be parametrized from the python config
+      tensorflow::setBackend(options, tensorflow::Backend::cpu);
 
       const std::string& entry_name = graph_entry.first;
       const std::string& graph_file = graph_entry.second;