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

Print tensor

Browse source
This commit is contained in:
lucasdelimanogueira 2024-04-29 21:09:50 -03:00
parent 1de187af93
commit f89292d3b5
1 changed files with 9 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

11
csrc/tensor.cpp
View file

@ -66,12 +66,19 @@ extern "C" {
}
float get_item(Tensor* tensor, int* indices) {
int index = 0;
for (int i = 0; i < tensor->ndim; i++) {
index += indices[i] * tensor->strides[i];
}
return tensor->data[index];
float result;
if (strcmp(tensor->device, "cuda") == 0) {
cudaMemcpy(&result, tensor->data + index, sizeof(float), cudaMemcpyDeviceToHost);
} else {
result = tensor->data[index];
}
return result;
}
void to_device(Tensor* tensor, char* target_device) {