Right but for such tasks, you want to do a very large number of small tasks.
So you won’t be able to run this on the GPU as it will be super slow.
And you will have to run this one by one on the CPU in a for-loop. Moving that for loop from python to cpp will speed it up but won’t give you "good’ perf I’m afraid…
You can use the cpp inline extensions to try this:
import torch
from torch.utils import cpp_extension
import time

cpp_source = """
std::vector<torch::Tensor> test_fn(std::vector<torch::Tensor> A, std::vector<torch::Tensor> B) {
    std::vector<torch::Tensor> result(A.size());

    for (int i=0; i<A.size(); ++i) {
        result[i] = A[i].mm(B[i]);
    }

    return result;
}
"""

def py_version(A, B):
    C = []
    for a, b in zip(A, B):
        C.append(a.mm(b))
    return C

mod = cpp_extension.load_inline(
    name="mod",
    cpp_sources=cpp_source,
    functions=["test_fn",],
)


for max_val in [10, 100, 1000, 1500]:
    print("")
    print("TESTING FOR {}".format(max_val))
    A = [torch.randn(i, i) for i in range(1, max_val)]
    B = [torch.randn(i, 3) for i in range(1, max_val)]

    t1 = time.time()
    res = mod.test_fn(A, B)
    cpp_time = time.time() - t1
    t1 = time.time()
    res = py_version(A, B)
    py_time = time.time() - t1

    print("cpp time: {}".format(cpp_time))
    print("py time: {}".format(py_time))