相关疑难解决方法(0)

我是 Numba 新手，我需要使用 Numba 来加速一些 Pytorch 功能。但我发现即使是一个非常简单的功能也不起作用:(

import torch
import numba

@numba.njit()
def vec_add_odd_pos(a, b):
    res = 0.
    for pos in range(len(a)):
        if pos % 2 == 0:
            res += a[pos] + b[pos]
    return res

x = torch.tensor([3, 4, 5.])
y = torch.tensor([-2, 0, 1.])
z = vec_add_odd_pos(x, y)

但出现以下错误

def vec_add_odd_pos(a, b):
    res = 0.
    ^

This error may have been caused by the following argument(s):
- argument 0: cannot determine Numba type of <class 'torch.Tensor'>
- argument 1: …