fix kinfer model (#124)

kscalelabs · Dec 5, 2024 · 74cb9d7 · 74cb9d7
1 parent 44dea1f
commit 74cb9d7
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/examples/gpr_walking.kinfer b/examples/gpr_walking.kinfer
diff --git a/examples/kinfer_policy.onnx b/examples/kinfer_policy.onnx
diff --git a/sim/sim2sim.py b/sim/sim2sim.py
@@ -20,7 +20,7 @@
 from tqdm import tqdm
 
 from sim.h5_logger import HDF5Logger
-from sim.model_export import ActorCfg, get_actor_policy
+from model_export import ActorCfg, get_actor_policy
 from kinfer.export.pytorch import export_to_onnx
 from kinfer.inference.python import ONNXModel
 
@@ -334,7 +334,7 @@ def run_mujoco(
             cycle_time=policy_cfg.cycle_time,
         )
 
-    if args.load_model.endswith(".onnx"):
+    if args.load_model.endswith(".kinfer"):
         policy = ONNXModel(args.load_model)
     else:
         actor_model, sim2sim_info, input_tensors = get_actor_policy(args.load_model, policy_cfg)
@@ -369,4 +369,4 @@ def run_mujoco(
         args.log_h5,
         args.render,
         args.h5_out_dir,
-    )
+    )