Hello! Today I compared numba.jit and torch.jit and was very surprised. What am I doing wrong?

image.png721×327 25.8 KB

import torch
from numba import jit

@torch.jit.script
def torch_jit_sum(x : torch.Tensor):
    res = torch.zeros_like(x[0, 0])
    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            res += x[i, j]
    return res

def blablabla(x):
    with torch.no_grad():
        return torch_jit_sum(x)

def loop_sum(x):
    res = torch.zeros_like(x[0, 0])
    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            res += x[i, j]
    return res

@jit
def numba_sum(x):
    res = 0.0
    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            res += x[i, j]
    return res