diff --git a/byol/main_loop.py b/byol/main_loop.py
index 544bbf2..fec16cf 100644
--- a/byol/main_loop.py
+++ b/byol/main_loop.py
@@ -145,7 +145,7 @@ def main(_):
                                     FLAGS.batch_size)
   else:
     raise ValueError(f'Unknown experiment mode: {FLAGS.experiment_mode}')
-  config['checkpointing_config']['checkpoint_dir'] = FLAGS.checkpoint_root
+  config['checkpointing_config']['checkpoint_dir'] = FLAGS.checkpoint_root  # pytype: disable=unsupported-operands  # dict-kwargs
 
   if FLAGS.worker_mode == 'train':
     train_loop(experiment_class, config)