Skip to content
This repository has been archived by the owner on Jul 1, 2024. It is now read-only.

Updated amg.py to allow batching of large mask sets and avoid over-sized torch.nonzero() calls #569

Open
wants to merge 8 commits into
base: main
Choose a base branch
from

Conversation

MartinKlefas
Copy link

@MartinKlefas MartinKlefas commented Sep 21, 2023

Fixes issue where Tensor.nonzero would fail on GPU for tensors containing more than INT_MAX elements.

Changes made:

  • Implemented a chunking mechanism to handle tensors in manageable sizes.
  • Each chunk's size is safely under the INT_MAX limit.
  • Non-zero indices from each chunk are concatenated to produce the final result.

This resolves the issue reported in #554 which correlates with the PyTorch issue pytorch/pytorch#51871.

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Sep 21, 2023
Copy link

@ron44-5 ron44-5 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Delete all permanently and remove all bug reset harder

@MartinKlefas
Copy link
Author

Delete all permanently and remove all bug reset harder

Hi, sorry if I've misunderstood something - but are these instructions for changes I need to make to the PR?

Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants