srush

import torch
from torch.autograd import Variable

q = Variable(torch.Tensor(5), requires_grad=True)
leaves = [Variable(torch.zeros(5, 5), requires_grad=True) for _ in range(10)]
intermediates = [l  for i, l in enumerate(leaves)]
loss = sum(q * v * i for i, v in enumerate(intermediates)).sum()

# define a helper for dividing intermediates into groups
def group(l, group_size):

srush

import torch
from torch.autograd import Variable

leaves = [Variable(torch.zeros(5, 5), requires_grad=True) for _ in range(10)]
intermediates = [l + i for i, l in enumerate(leaves)]
loss = sum(v * i for i, v in enumerate(intermediates)).sum()

# define a helper for dividing intermediates into groups
def group(l, group_size):
    """Groups l into chunks of size group_size.