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

Set tensors to cpu in _append_to_hooked_list #278

Closed
danbraunai-apollo opened this issue Jan 12, 2024 · 0 comments · Fixed by #355
Closed

Set tensors to cpu in _append_to_hooked_list #278

danbraunai-apollo opened this issue Jan 12, 2024 · 0 comments · Fixed by #355
Labels
enhancement New feature or request priority-medium

Comments

@danbraunai-apollo
Copy link
Contributor

Currently, rib.hook_fns._append_to_hooked_list looks like:

def _append_to_hooked_list(
    hooked_data: dict[str, Any],
    hook_name: str,
    data_key: str,
    element_to_append: Any,
) -> None:
    """Append the given element to a hooked list. Creates the list if it doesn't exist.

    Args:
        hooked_data: Dictionary of hook data that will be updated.
        hook_name: Name of hook. Used as a 1st-level key in `hooked_data`.
        data_key: Name of data. Used as a 2nd-level key in `hooked_data`.
        element_to_append: Appended to hooked data.
    """
    hooked_data.setdefault(hook_name, {}).setdefault(data_key, [])
    hooked_data[hook_name][data_key].append(element_to_append)

We should convert to CPU when storing these, and handle converting back to GPU later wherever these things tensors are used.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request priority-medium
Projects
None yet
Development

Successfully merging a pull request may close this issue.

1 participant