set input tensor dtype

commaai · Sep 18, 2024 · b5e822a · b5e822a
1 parent f51aa0f
commit b5e822a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/openpilot/compile2.py b/openpilot/compile2.py
@@ -30,9 +30,10 @@ def get_schedule(onnx_data) -> Tuple[List[ScheduleItem], List[ScheduleItem]]:
   onnx_model = onnx.load(io.BytesIO(onnx_data))
   run_onnx = get_run_onnx(onnx_model)
   input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
+  input_types = {inp.name:onnx.helper.tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type) for inp in onnx_model.graph.input}
 
   # run the model
-  inputs = {k:Tensor.empty(*shp) for k,shp in input_shapes.items()}
+  inputs = {k:Tensor.empty(*shp, dtype=getattr(dtypes, input_types[k].name)) for k,shp in input_shapes.items()}
   ret: Tensor = next(iter(run_onnx(inputs).values())).cast(dtypes.float32).contiguous()
   schedule = ret.lazydata.schedule()