diff --git a/examples/forward_problems/mechanics/example_hyper_visco_2d.py b/examples/forward_problems/mechanics/example_hyper_visco_2d.py
index cc17b01..627c6bd 100644
--- a/examples/forward_problems/mechanics/example_hyper_visco_2d.py
+++ b/examples/forward_problems/mechanics/example_hyper_visco_2d.py
@@ -94,7 +94,7 @@ def dirichlet_bc_func(xs, t, nn):
 ##################
 opt = Adam(loss_function, learning_rate=1.0e-3, has_aux=True, clip_gradients=False)
 opt, opt_st = opt.init(params)
-for epoch in range(100000):
+for epoch in range(1000):
     params, opt_st, loss = opt.step(params, opt_st, problem)
     # logger.log_loss(loss, epoch)
     if epoch % 100 == 0:
diff --git a/examples/forward_problems/mechanics/example_hyper_visco_3d.py b/examples/forward_problems/mechanics/example_hyper_visco_3d.py
index 657eba5..6f55bd7 100644
--- a/examples/forward_problems/mechanics/example_hyper_visco_3d.py
+++ b/examples/forward_problems/mechanics/example_hyper_visco_3d.py
@@ -3,7 +3,7 @@
 ##################
 # for reproducibility
 ##################
-key = random.key(10)
+key = random.PRNGKey(10)
 
 ##################
 # file management
@@ -93,9 +93,9 @@ def dirichlet_bc_func(xs, t, nn):
 # train network
 ##################
 opt = Adam(loss_function, learning_rate=1.0e-3, has_aux=True, clip_gradients=False)
-opt_st = opt.init(params)
+opt, opt_st = opt.init(params)
 for epoch in range(100000):
-    params, opt_st, loss = opt.step(params, problem, opt_st)
+    params, opt_st, loss = opt.step(params, opt_st, problem)
     # logger.log_loss(loss, epoch)
     if epoch % 10 == 0:
         print(epoch, flush=True)
@@ -106,6 +106,7 @@ def dirichlet_bc_func(xs, t, nn):
     if epoch % 1000 == 0:
         print("writing exodus output")
         pp.init(
+            params,
             problem,
             f"output_{str(epoch).zfill(6)}.e",
             node_variables=["field_values"],
diff --git a/examples/forward_problems/mechanics/example_incompressible_2d.py b/examples/forward_problems/mechanics/example_incompressible_2d.py
index 4382695..546e06a 100644
--- a/examples/forward_problems/mechanics/example_incompressible_2d.py
+++ b/examples/forward_problems/mechanics/example_incompressible_2d.py
@@ -4,7 +4,7 @@
 # for reproducibility
 ##################
 key = random.PRNGKey(10)
-key = random.split(key, 8) # comment this to not use an ensemble
+# key = random.split(key, 8) # comment this to not use an ensemble
 
 ##################
 # file management
@@ -72,7 +72,7 @@
 pp.init(params, problem, 'output.e', 
   node_variables=[
     'field_values',
-    # 'internal_force'
+    'internal_force'
   ], 
   element_variables=[
     # 'deformation_gradient',