Cleaned up code

Author: ajkv-google

recml/core/training/partitioning.py

@@ -22,7 +22,6 @@
 import flax.linen as nn
 import jax
 import numpy as np
-# FIXED: Use the public experimental module available in JAX 0.4.30
 from recml.core.training import mesh_context
 
 

recml/examples/train_hstu_keras.py

@@ -15,7 +15,6 @@
 # Add the RecML folder to the system path
 sys.path.append(os.path.join(os.getcwd(), "../../../RecML"))
 
-# RecML Imports
 from recml.core.training import core
 from recml.core.training import keras_trainer
 from recml.layers.keras import hstu