fix

mlcommons · Dec 7, 2023 · 3439204 · 3439204
1 parent 6c25b2b
commit 3439204
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/algorithmic_efficiency/workloads/fastmri/fastmri_jax/workload.py b/algorithmic_efficiency/workloads/fastmri/fastmri_jax/workload.py
@@ -41,7 +41,7 @@ def init_model_fn(
     console_kwargs={
         'force_terminal': False, 'force_jupyter': False, 'width': 240},
     )
-    print(tabulate_fn(fake_inputs, train=False))
+    print(tabulate_fn(fake_batch, train=False))
     variables = jax.jit(self._model.init)({'params': rng}, fake_batch)
     params = variables['params']
     self._param_shapes = param_utils.jax_param_shapes(params)