training = True causes a GPU memory error, asks for 53 GB on normal sized proteins

Hello,

Sorry for bothering again,

I've been looking into manipulating the function parameters of the batch run function in the colab, This one, and got a crash.

Trying to isolate the parameters I've changed, it seems just activating training = True is enough, here are the parameters, crash message, and the input sequence :

---- parameters

all default params except :

• input and output dirs, but it finds the right files and seems to read them fine (output the right pdb names)
• stop at score 90 (nevermind, tried the default 100 too)
• training = True

---- crash message

2023-01-11 09:30:23,449 More than one sequence in /content/drive/MyDrive/ColabFold_files/batch/input_fasta/rcsb_pdb_1JEV.fasta, ignoring all but the first sequence
2023-01-11 09:30:23,452 More than one sequence in /content/drive/MyDrive/ColabFold_files/batch/input_fasta/rcsb_pdb_5KZT.fasta, ignoring all but the first sequence
2023-01-11 09:30:23,454 More than one sequence in /content/drive/MyDrive/ColabFold_files/batch/input_fasta/rcsb_pdb_7JLS.fasta, ignoring all but the first sequence
2023-01-11 09:30:23,465 Found 5 citations for tools or databases
2023-01-11 09:30:27,958 Query 1/3: rcsb_pdb_1JEV (length 517)
2023-01-11 09:30:29,785 Running model_3
2023-01-11 09:30:52,301 Could not predict rcsb_pdb_1JEV. Not Enough GPU memory? RESOURCE_EXHAUSTED: Out of memory while trying to allocate 53061779456 bytes.
2023-01-11 09:30:52,329 Query 2/3: rcsb_pdb_5KZT (length 536)
2023-01-11 09:30:53,880 Running model_3
2023-01-11 09:31:15,591 Could not predict rcsb_pdb_5KZT. Not Enough GPU memory? RESOURCE_EXHAUSTED: Out of memory while trying to allocate 53061779456 bytes.
2023-01-11 09:31:15,622 Query 3/3: rcsb_pdb_7JLS (length 553)
2023-01-11 09:31:16,237 Running model_3
2023-01-11 09:31:32,319 Could not predict rcsb_pdb_7JLS. Not Enough GPU memory? RESOURCE_EXHAUSTED: Out of memory while trying to allocate 53061779456 bytes.
2023-01-11 09:31:32,348 Done

---- example of fasta file

>7JLS_1|Chain A|Probable periplasmic dipeptide-binding lipoprotein DppA|Mycobacterium tuberculosis (strain ATCC 25618 / H37Rv) (83332)
MAGLNDIFEAQKIEWHELEVLFQGPMSPDVVLVNGGEPPNPLIPTGTNDSNGGRIIDRLFAGLMSYDAVGKPSLEVAQSIESADNVNYRITVKPGWKFTDGSPVTAHSFVDAWNYGALSTNAQLQQHFFSPIEGFDDVAGAPGDKSRTTMSGLRVVNDLEFTVRLKAPTIDFTLRLGHSSFYPLPDSAFRDMAAFGRNPIGNGPYKLADGPAGPAWEHNVRIDLVPNPDYHGNRKPRNKGLRFEFYANLDTAYADLLSGNLDVLDTIPPSALTVYQRDLGDHATSGPAAINQTLDTPLRLPHFGGEEGRLRRLALSAAINRPQICQQIFAGTRSPARDFTARSLPGFDPNLPGNEVLDYDPQRARRLWAQADAISPWSGRYAIAYNADAGHRDWVDAVANSIKNVLGIDAVAAPQPTFAGFRTQITNRAIDSAFRAGWRGDYPSMIEFLAPLFTAGAGSNDVGYINPEFDAALAAAEAAPTLTESHELVNDAQRILFHDMPVVPLWDYISVVGWSSQVSNVTVTWNGLPDYENIVKAENLYFQGGHHHHHHHH
>7JLS_2|Chain B|Peptide SER-VAL-ALA|Escherichia coli BL21(DE3) (469008)
SVA

There are two chains but it warns that it only takes a single one (batch for multimer is another thing I have to look into), so it should run as monomer fine, especially with proteins of size 500 to 550.

The first thing I can think of that could do this was if the flag training = True caused to compute all the gradients for backprop, but by my understanding it's supposed to only activate the dropout layers ?

Start with the batch/AlphaFold2_batch.ipynb notebook linked in the report and reproduce the run using training=True with the 500–550 residue FASTA inputs. Compare this path with the default parameters and inspect the reported 53 GB allocation failure. Done means these normal-sized monomer inputs no longer fail unexpectedly when training=True, with the behavior verified against the reported case.