diff --git a/ogb_lsc/mag/experiment.py b/ogb_lsc/mag/experiment.py
index 1dfcbfd..7f9e610 100644
--- a/ogb_lsc/mag/experiment.py
+++ b/ogb_lsc/mag/experiment.py
@@ -154,8 +154,8 @@ class Experiment(experiment.AbstractExperiment):
   ):
     """Initializes experiment."""
     super(Experiment, self).__init__(mode=mode, init_rng=init_rng)
-    tf.config.experimental.set_visible_devices([], device_type='gpu')
-    tf.config.experimental.set_visible_devices([], device_type='tpu')
+    tf.config.experimental.set_visible_devices([], device_type='GPU')
+    tf.config.experimental.set_visible_devices([], device_type='TPU')
 
     if mode not in ('train', 'eval', 'train_eval_multithreaded'):
       raise ValueError(f'Invalid mode {mode}.')
diff --git a/ogb_lsc/pcq/experiment.py b/ogb_lsc/pcq/experiment.py
index 522798f..5e66a86 100644
--- a/ogb_lsc/pcq/experiment.py
+++ b/ogb_lsc/pcq/experiment.py
@@ -90,8 +90,8 @@ class Experiment(experiment.AbstractExperiment):
       raise ValueError(f'Invalid mode {mode}.')
 
     # Do not use accelerators in data pipeline.
-    tf.config.experimental.set_visible_devices([], device_type='gpu')
-    tf.config.experimental.set_visible_devices([], device_type='tpu')
+    tf.config.experimental.set_visible_devices([], device_type='GPU')
+    tf.config.experimental.set_visible_devices([], device_type='TPU')
 
     self.mode = mode
     self.init_rng = init_rng