Fixed tensor freezing indexing issue

Author: DoeringChristian

src/python/freeze.cpp

@@ -681,13 +681,13 @@ void FlatVariables::traverse(nb::handle h, TraverseContext &ctx) {
 
                 auto full_shape = nb::borrow<nb::tuple>(shape(h));
 
-                nb::list outer_shape;
+                nb::list inner_shape;
                 if (full_shape.size() > 0)
-                    for (uint32_t i = 0; i < full_shape.size() - 1; i++) {
-                        outer_shape.append(full_shape[i]);
+                    for (uint32_t i = 1; i < full_shape.size(); i++) {
+                        inner_shape.append(full_shape[i]);
                     }
 
-                layout.py_object = nb::tuple(outer_shape);
+                layout.py_object = nb::tuple(inner_shape);
 
                 traverse(nb::steal(array), ctx);
             } else if (s.ndim != 1) {
@@ -825,15 +825,15 @@ nb::object FlatVariables::construct() {
             if (s.is_tensor) {
                 nb::object array = construct();
 
-                auto outer_shape = nb::borrow<nb::tuple>(layout.py_object);
+                auto inner_shape = nb::borrow<nb::tuple>(layout.py_object);
                 auto last_dim    = prod(shape(array), nb::none())
-                                    .floor_div(prod(outer_shape, nb::none()));
+                                    .floor_div(prod(inner_shape, nb::none()));
 
                 nb::list full_shape;
-                for (uint32_t i = 0; i < outer_shape.size(); i++) {
-                    full_shape.append(outer_shape[i]);
-                }
                 full_shape.append(last_dim);
+                for (uint32_t i = 0; i < inner_shape.size(); i++) {
+                    full_shape.append(inner_shape[i]);
+                }
 
                 nb::object tensor = layout.type(array, nb::tuple(full_shape));
                 return tensor;