testcase.py

import torch
from functorch import grad

device='cpu'
x = torch.randn([], requires_grad=True)
a = torch.tensor([1], device=device)
b = torch.tensor([1, 2, 3], device=device)
c = torch.tensor([1, 2], device=device)

def foo(x):
    y = torch.tensor(list(product([a], b, c)), device=device)
    return (x + y).sum()

grad(foo)(x)