Updated data generation code

model.py

@@ -46,7 +46,7 @@ def init_fn(key: jax.Array, config: ModelConfig):
 
 def forward(params: dict, input_timesteps: jax.Array, config: ModelConfig):
     """
-    Model's forward function. Takes in the parameters and inptu timesteps, returns predictions
+    Model's forward function. Takes in the parameters and input timesteps, returns predictions
     """
     batch_size, num_agents, _ = input_timesteps.shape
 
@@ -119,6 +119,7 @@ def train_model(config: ModelConfig, inputs: jax.Array, targets: jax.Array,
             params, opt_state, loss_val = update_step(params, opt_state, x, y, config)
             running_loss += loss_val
 
+
         epoch_loss = running_loss / num_batches
         loss_history[f"epoch_{epoch}"].append(epoch_loss)