From 493f6b0b8345ea20714c199905003a82a7f320f0 Mon Sep 17 00:00:00 2001
From: lucasdelimanogueira <lucasln96@hotmail.com>
Date: Mon, 29 Apr 2024 22:22:20 -0300
Subject: [PATCH] Fix set device

---
 csrc/tensor.cpp | 17 +++++++++++++++--
 norch/tensor.py |  7 +++++--
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/csrc/tensor.cpp b/csrc/tensor.cpp
index a46a9f4..03b2b6a 100644
--- a/csrc/tensor.cpp
+++ b/csrc/tensor.cpp
@@ -20,7 +20,14 @@ extern "C" {
         tensor->data = data;
         tensor->shape = shape;
         tensor->ndim = ndim;
-        tensor->device = device;
+        
+        tensor->device = (char*)malloc(strlen(device) + 1);
+        if (device != NULL) {
+            strcpy(tensor->device, device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
 
         tensor->size = 1;
         for (int i = 0; i < ndim; i++) {
@@ -105,7 +112,13 @@ extern "C" {
             exit(1);
         }
 
-        char* device = tensor1->device;
+        char* device = (char*)malloc(strlen(tensor1->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor1->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
         int ndim = tensor1->ndim;
         int* shape = (int*)malloc(ndim * sizeof(int));
         if (shape == NULL) {
diff --git a/norch/tensor.py b/norch/tensor.py
index 3c2a66e..982e202 100644
--- a/norch/tensor.py
+++ b/norch/tensor.py
@@ -23,7 +23,6 @@ class Tensor:
             self.ndim_ctype = ctypes.c_int(len(shape))
             self.device_ctype = device.encode('utf-8')
 
-
             self.shape = shape
             self.ndim = len(shape)
             self.device = device
@@ -121,7 +120,7 @@ class Tensor:
         index = [0] * self.ndim
         result = "tensor(["
         result += print_recursively(self, 0, index)
-        result += "])"
+        result += f"""], device="{self.device}")"""
         return result
 
     def __repr__(self):
@@ -140,6 +139,7 @@ class Tensor:
         result_data.tensor = result_tensor_ptr
         result_data.shape = self.shape.copy()
         result_data.ndim = self.ndim
+        result_data.device = self.device
 
         return result_data
     
@@ -156,6 +156,7 @@ class Tensor:
         result_data.tensor = result_tensor_ptr
         result_data.shape = self.shape.copy()
         result_data.ndim = self.ndim
+        result_data.device = self.device
 
         return result_data
     
@@ -172,6 +173,7 @@ class Tensor:
         result_data.tensor = result_tensor_ptr
         result_data.shape = self.shape.copy()
         result_data.ndim = self.ndim
+        result_data.device = self.device
 
         return result_data
     
@@ -191,6 +193,7 @@ class Tensor:
         result_data.tensor = result_tensor_ptr
         result_data.shape = [self.shape[0], other.shape[1]]
         result_data.ndim = 2
+        result_data.device = self.device
 
         return result_data