a few improvements

Author: LeoGrin

src/tabpfn/misc/onnx_wrapper.py

@@ -249,7 +249,7 @@ def export_model(
 
         # Define dynamic axes for variable input sizes
         dynamic_axes = {
-            "X": {0: "num_datapoints", 1: "batch_size", 2: "num_features"},
+            "X": {0: "num_datapoints", 2: "num_features"},
             "y": {0: "num_labels"},
             "single_eval_pos": {},
             "only_return_standard_out": {},
@@ -291,9 +291,10 @@ def check_input_names(model_path: str) -> None:
         model_path: The path to the ONNX model file.
     """
     onnx.load(model_path)
-    # get input names from graph
-    graph = onnx.load(model_path).graph
-    [input_node.name for input_node in graph.input]
+
+    # Print input names
+
+    # Print output names
 
 
 if __name__ == "__main__":

src/tabpfn/model/encoders.py

@@ -90,7 +90,7 @@ def normalize_data(
             mean = torch_nanmean(data, axis=0)  # type: ignore
             std = torch_nanstd(data, axis=0) + 1e-20
 
-        if len(data) == 1 or normalize_positions == 1:
+        if data.shape[0] == 1 or normalize_positions == 1:
             std[:] = 1.0
 
         if std_only: