Minor changes to main.py

saforem2 · Apr 14, 2022 · db3b726 · db3b726
1 parent 465da8a
commit db3b726
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/l2hmc/main.py b/src/l2hmc/main.py
@@ -18,7 +18,6 @@
 def train_tensorflow(cfg: DictConfig) -> dict:
     import tensorflow as tf
     tf.keras.backend.set_floatx(cfg.precision)
-    # tf.keras.backend.set_floatx('float32') # or 'float64 for double precision
     # assert tf.keras.backend.floatx() == tf.float32
     import horovod.tensorflow as hvd
     hvd.init()
@@ -64,6 +63,7 @@ def main(cfg: DictConfig) -> None:
 
     if framework in ['tf', 'tensorflow']:
         _ = train_tensorflow(cfg)
+
     elif framework in ['pt', 'pytorch']:
         _ = train_pytorch(cfg)