getting it to level 0

Author: init-22

algoperf/workloads/criteo1tb/criteo1tb_jax/workload.py

@@ -13,7 +13,7 @@
 from algoperf.workloads.criteo1tb.criteo1tb_jax import models
 from algoperf.workloads.criteo1tb.workload import \
     BaseCriteo1TbDlrmSmallWorkload
-from custom_pytorch_jax_converter import use_pytorch_weights_inplace, use_pytorch_weights_inplace_mnist
+from custom_pytorch_jax_converter import use_pytorch_weights_inplace
 
 
 
@@ -108,7 +108,6 @@ def init_model_fn(
     initial_params = use_pytorch_weights_inplace(initial_params, file_name="/results/pytorch_base_model_criteo1tb_24_june.pth")
     self._param_shapes = param_utils.jax_param_shapes(initial_params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)
-    return initial_params, None
     return jax_utils.replicate(initial_params), None
 
   def is_output_params(self, param_key: spec.ParameterKey) -> bool:

reference_algorithms/schedule_free/jax/submission.py

@@ -172,8 +172,8 @@ def update_params(workload: spec.Workload,
 
   # Log the number of parameters.
   if global_step % 100 == 0 and workload.metrics_logger is not None:
-    date_ = "2025=06-14"
-    file_name = f"/results/schedule_free_test_pytorch_weights/criteo1tb_{date_}_after_{global_step}.pth"
+    date_ = "2025-06-14"
+    file_name = f"/results/schedule_free_test_pytorch_weights/criteo1tb_{date_}_after_{global_step}_steps.pth"
     params = use_pytorch_weights2(new_params, file_name=file_name, replicate=True)
     are_weights_equal(new_params, params)
     del params