f41ae3bce45a4cf2f101d610704838d335b1f586,tensorly/backend/pytorch_backend.py,,where,#Any#Any#Any#,101

Before Change

    out_ = out.view(N)

    for i in range(N):
        out_[i] = x_[i] if condition_[i] else y_[i]
    return out

After Change

def where(condition, x, y):
    assert condition.shape == x.shape == y.shape, "Dimension mismatch"
    result = zeros_like(condition, **context(x))
    result[condition] = x
    result[~condition] = y
    return result

Instances