setting up to run the workloads

Author: init-22

algoperf/workloads/criteo1tb/criteo1tb_jax/workload.py

@@ -11,7 +11,6 @@
 from algoperf import param_utils, spec
 from algoperf.workloads.criteo1tb.criteo1tb_jax import models
 from algoperf.workloads.criteo1tb.workload import BaseCriteo1TbDlrmSmallWorkload
-from custom_pytorch_jax_converter import use_pytorch_weights
 
 
 class Criteo1TbDlrmSmallWorkload(BaseCriteo1TbDlrmSmallWorkload):
@@ -105,7 +104,6 @@ def init_model_fn(
       jnp.ones(input_shape, jnp.float32),
     )
     initial_params = initial_variables['params']
-    initial_params = use_pytorch_weights(file_name="~/results/pytorch_base_model_criteo1tb_1_july.pth")
     self._param_shapes = param_utils.jax_param_shapes(initial_params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)
     return jax_utils.replicate(initial_params), None

algoperf/workloads/criteo1tb/criteo1tb_pytorch/workload.py

@@ -82,7 +82,6 @@ def init_model_fn(self, rng: spec.RandomState) -> spec.ModelInitState:
       use_layer_norm=self.use_layer_norm,
       embedding_init_multiplier=self.embedding_init_multiplier,
     )
-    torch.save(model.state_dict(), '~/results/pytorch_base_model_criteo1tb_1_july.pth')
     self._param_shapes = param_utils.pytorch_param_shapes(model)
     self._param_types = param_utils.pytorch_param_types(self._param_shapes)
     model.to(DEVICE)