diff --git a/python/__pycache__/tensor.cpython-39.pyc b/python/__pycache__/tensor.cpython-39.pyc
index cfb682e..53e9673 100644
Binary files a/python/__pycache__/tensor.cpython-39.pyc and b/python/__pycache__/tensor.cpython-39.pyc differ
diff --git a/python/tensor.py b/python/tensor.py
index fb468b9..e9a7d0d 100644
--- a/python/tensor.py
+++ b/python/tensor.py
@@ -9,7 +9,9 @@ class CTensor(ctypes.Structure):
     ]
 
 class Tensor:
-    def __init__(self, data, shape):
+    def __init__(self, data):
+        data, shape = self.flatten(data)
+        print(data, shape)
         self.lib = ctypes.CDLL("../build/libtensor.so")  # Adjust the path to the shared library
         self.data = (ctypes.c_float * len(data))(*data)
         self.shape = shape
@@ -23,6 +25,14 @@ class Tensor:
             (ctypes.c_int * len(shape))(*shape),
             ctypes.c_int(len(shape))
         )
+
+    def flatten(self, nested_list):
+        flat_data = []
+        shape = [len(nested_list), len(nested_list[0])]
+        for sublist in nested_list:
+            for item in sublist:
+                flat_data.append(item)
+        return flat_data, shape
     
     def __getitem__(self, indices):
         if len(indices) != self.ndim:
diff --git a/python/tests.ipynb b/python/tests.ipynb
index f56bde6..ca51d7d 100644
--- a/python/tests.ipynb
+++ b/python/tests.ipynb
@@ -9,36 +9,23 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "kkkk <tensor.c_int_Array_2 object at 0x7f990806fcc0> @@ c_int(2)\n",
+      "[0, 1, 2, 0, 2, 2] [2, 3, 3]\n",
+      "Tensor shape: [2, 3, 3]\n",
       "Tensor created successfully\n",
       "Tensor information:\n",
-      "Number of dimensions: 2\n",
-      "Shape: [3, 3]\n",
+      "Number of dimensions: 3\n",
+      "Shape: [2, 3, 3]\n",
       "Data:\n",
-      "[0.10, 0.20, 0.30, 0.40, 0.50, 0.60, 0.70, 0.80, 0.90]\n"
-     ]
-    },
-    {
-     "ename": "ValueError",
-     "evalue": "Number of indices must match the number of dimensions",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[1], line 6\u001b[0m\n\u001b[1;32m      4\u001b[0m shape \u001b[38;5;241m=\u001b[39m [\u001b[38;5;241m3\u001b[39m, \u001b[38;5;241m3\u001b[39m]\n\u001b[1;32m      5\u001b[0m tensor \u001b[38;5;241m=\u001b[39m Tensor(data, shape)\n\u001b[0;32m----> 6\u001b[0m \u001b[43mTensor\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__getitem__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mtensor\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m~/Documentos/recreate_pytorch/foo/python/tensor.py:15\u001b[0m, in \u001b[0;36mTensor.__getitem__\u001b[0;34m(self, indices)\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__getitem__\u001b[39m(\u001b[38;5;28mself\u001b[39m, indices):\n\u001b[1;32m     14\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(indices) \u001b[38;5;241m!=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mndim:\n\u001b[0;32m---> 15\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mNumber of indices must match the number of dimensions\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     16\u001b[0m     indices \u001b[38;5;241m=\u001b[39m (ctypes\u001b[38;5;241m.\u001b[39mc_int \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mlen\u001b[39m(indices))(\u001b[38;5;241m*\u001b[39mindices)\n\u001b[1;32m     17\u001b[0m     value \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlib\u001b[38;5;241m.\u001b[39mget_item(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtensor, indices)  \n",
-      "\u001b[0;31mValueError\u001b[0m: Number of indices must match the number of dimensions"
+      "[0.00, 1.00, 2.00, 0.00, 2.00, 2.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00]\n"
      ]
     }
    ],
    "source": [
     "from tensor import Tensor\n",
     "\n",
-    "data = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]\n",
-    "shape = [3, 3]\n",
-    "tensor = Tensor(data, shape)\n",
-    "tensor.shape\n",
-    "Tensor.__getitem__(tensor, [0, 0])"
+    "data = [[0, 1, 2], [0, 2, 2], [0, 1, 2]]\n",
+    "tensor = Tensor(data)\n",
+    "print(\"Tensor shape:\", tensor.shape)"
    ]
   },
   {
@@ -50,12 +37,33 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "-648269760\n"
+      "Tensor created successfully\n",
+      "Tensor information:\n",
+      "Number of dimensions: 2\n",
+      "Shape: [3, 3]\n",
+      "Data:\n",
+      "[0.10, 0.20, 0.30, 0.40, 0.50, 0.60, 0.70, 0.80, 0.90]\n"
      ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "0.10000000149011612 0.20000000298023224 0.30000001192092896 \n",
+       "0.4000000059604645 0.5 0.6000000238418579 \n",
+       "0.699999988079071 0.800000011920929 0.8999999761581421"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
     }
    ],
    "source": [
-    "print(tensor[1, 1])  # Accessing an element"
+    "\n",
+    "data = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]\n",
+    "shape = [3, 3]\n",
+    "tensor = Tensor(data, shape)\n",
+    "tensor\n"
    ]
   },
   {