aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix set device

Browse source
This commit is contained in:
lucasdelimanogueira 2024-04-29 22:22:20 -03:00
parent f89292d3b5
commit 493f6b0b83
2 changed files with 20 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

17
csrc/tensor.cpp
View file

@ -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) {

7
norch/tensor.py
View file

@ -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