From 3174cf9af51913c08e302f56dd0a10648cb6212c Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Tue, 30 Apr 2024 02:19:17 -0300 Subject: [PATCH] improve cpu.h header --- csrc/cpu.h | 9 ++++++++- test.py | 2 +- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/csrc/cpu.h b/csrc/cpu.h index 4370dc8..3f76c23 100644 --- a/csrc/cpu.h +++ b/csrc/cpu.h @@ -1,5 +1,12 @@ +#ifndef CPU_H +#define CPU_H + +#include "tensor.h" + void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); -void pow_tensor_cpu(Tensor* tensor, float power, float* result_data); \ No newline at end of file +void pow_tensor_cpu(Tensor* tensor, float power, float* result_data); + +#endif /* CPU_H */ diff --git a/test.py b/test.py index 62959be..bd2c647 100644 --- a/test.py +++ b/test.py @@ -16,7 +16,7 @@ def matrix_sum(matrix1, matrix2): if __name__ == "__main__": import norch - a = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]]) + a = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]]).to("cuda") b = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]]) import time import random