Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make to_device support list inputs #640

Open
KevinMusgrave opened this issue Jun 18, 2023 · 0 comments · May be fixed by #643
Open

Make to_device support list inputs #640

KevinMusgrave opened this issue Jun 18, 2023 · 0 comments · May be fixed by #643
Labels
enhancement New feature or request
Milestone

Comments

@KevinMusgrave
Copy link
Owner

KevinMusgrave commented Jun 18, 2023

As suggested and implemented in pull request #635 by @domenicoMuscill0.

I'm removing this feature from that PR because it's already a large PR.

Here's the implementation by @domenicoMuscill0:

def to_device(
    x: Union[torch.Tensor, nn.Parameter, List, Tuple],
    tensor=None,
    device=None,
    dtype: Union[torch.dtype, List, Tuple] = None,
):
    dv = device if device is not None else tensor.device
    is_iterable = is_list_or_tuple(x)
    if not is_iterable:
        x = [x]

    xd = x
    if is_list_or_tuple(dtype):
        if len(dtype) == len(x):
            xd = [
                to_dtype(x[i].to(dv), tensor=tensor, dtype=dtype[i])
                for i in range(len(x))
            ]
        else:
            raise RuntimeError(
                f"The size of dtype was {len(dtype)}. It is only available 1 or the same of x"
            )
    elif dtype is not None:
        xd = [to_dtype(xt.to(dv), tensor=tensor, dtype=dtype) for xt in x]

    if len(xd) == 1:
        xd = xd[0]
    return xd
@KevinMusgrave KevinMusgrave added the enhancement New feature or request label Jun 18, 2023
@KevinMusgrave KevinMusgrave added this to the v3.0 milestone Jul 3, 2023
@KevinMusgrave KevinMusgrave linked a pull request Jul 3, 2023 that will close this issue
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging a pull request may close this issue.

1 participant