d6e8ea4250b2a48262bae0276047371ea2dcccf1,onmt/utils/loss.py,,shards,#Any#Any#Any#,246

Before Change

            yield dict(zip(keys, shard_tensors))

        // Assumed backprop"d
        variables = ((state[k], v.grad.data) for k, v in non_none.items()
                     if isinstance(v, Variable) and v.grad is not None)
        inputs, grads = zip(*variables)
        torch.autograd.backward(inputs, grads)

After Change

        // want a sequence of dictionaries of tensors.
        // First, unzip the dictionary into a sequence of keys and a
        // sequence of tensor-like sequences.
        keys, values = zip(*((k, [v_chunk for v_chunk in v_split])
                             for k, (_, v_split) in non_none.items()))

        // Now, yield a dictionary for each shard. The keys are always
        // the same. values is a sequence of length //keys where each
        // element is a sequence of length //shards. We want to iterate
        // over the shards, not over the keys: therefore, the values need
        // to be re-zipped by shard and then each shard can be paired
        // with the keys.
        for shard_tensors in zip(*values):
            yield dict(zip(keys, shard_tensors))

        // Assumed backprop"d
        variables = []
        for k, (v, v_split) in non_none.items():
            if isinstance(v, torch.Tensor) and state[k].requires_grad:
                variables.extend(zip(torch.split(state[k], shard_size),
                                     [v_chunk.grad for v_chunk in v_split]))
        inputs, grads = zip(*variables)
        torch.autograd.backward(inputs, grads)

Instances