Gustaf Ahdritz

Hm so that's not it either. Whatever the issue is, it's probably happening somewhere in that function. Would you mind pinpointing where the samples are disappearing?

In the reroll() function, could you verify that `len(self._samples) == 1` and also that `torch.sum(dataset_choices) == 0`?

To clarify, `self._samples` contains one infinite iterator for each dataset you're using. Each entry thereof feeds you samples from the corresponding dataset shuffled + filtered in various ways forever. Yes,...

Could you verify programmatically that every single chain in your `alignment_dir` has a corresponding .mmcif file in the `data_dir`? Take all chain names in the former, split on `_`, and...

Delete their `alignment_dirs` and rerun. I'll look into what's causing this.

 My `chain_data_cache` has that chain, and I tried parsing my copy of `6tif.cif` and found them there too. Could you verify that [this file](https://drive.google.com/file/d/18tU1txzDrYug0CBs4S16kEFpd9UM15c6/view?usp=sharing) does not match the copy...

Are you saying that it didn't OOM with the full 48 blocks?

Is this the 256 or 384 setting?

Does torch allocate more total memory when you run it with 48 blocks? Sometimes it seems like torch opportunistically allocates less than it ultimately ends up needing. My hunch is...

All I can say on this for now is that we're working on more memory-efficient attention. In principle, there's no reason why we shouldn't be able to get it as...