Out of GPU memory on A100 80GB

https://www.rbvi.ucsf.edu/chimerax/data/alphafold-jan2022/afspeed.html

I was trying to predict a 3200 a.a. protein on our cluster with 2 A100 80GB GPU. Base on the page above, 48 GB seems sufficient for my protein length. However, I keep getting the out of GPU memory error as below. Can anyone help me with this?

2023-04-02 10:20:17.390936: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2040] Execution of replica 0 failed: Resource exhausted: Out of memory while trying to allocate 64793134928 bytes.
Traceback (most recent call last):
File "/app/alphafold/run_alphafold.py", line 422, in
app.run(main)
File "/opt/alphafoldenv/lib/python3.8/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/opt/alphafoldenv/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "/app/alphafold/run_alphafold.py", line 398, in main
predict_structure(
File "/app/alphafold/run_alphafold.py", line 198, in predict_structure
prediction_result = model_runner.predict(processed_feature_dict,
File "/app/alphafold/alphafold/model/model.py", line 167, in predict
result = self.apply(self.params, jax.random.PRNGKey(random_seed), feat)
File "/opt/alphafoldenv/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 183, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/opt/alphafoldenv/lib/python3.8/site-packages/jax/_src/api.py", line 424, in cache_miss
out_flat = xla.xla_call(
File "/opt/alphafoldenv/lib/python3.8/site-packages/jax/core.py", line 1560, in bind
return call_bind(self, fun, *args, **params)
File "/opt/alphafoldenv/lib/python3.8/site-packages/jax/core.py", line 1551, in call_bind
outs = primitive.process(top_trace, fun, tracers, params)
File "/opt/alphafoldenv/lib/python3.8/site-packages/jax/core.py", line 1563, in process
return trace.process_call(self, fun, tracers, params)
File "/opt/alphafoldenv/lib/python3.8/site-packages/jax/core.py", line 606, in process_call
return primitive.impl(f, *tracers, **params)
File "/opt/alphafoldenv/lib/python3.8/site-packages/jax/interpreters/xla.py", line 595, in _xla_call_impl
return compiled_fun(*args)
File "/opt/alphafoldenv/lib/python3.8/site-packages/jax/interpreters/xla.py", line 893, in _execute_compiled
out_bufs = compiled.execute(input_bufs)
jax._src.traceback_util.UnfilteredStackTrace: RuntimeError: Resource exhausted: Out of memory while trying to allocate 64793134928 bytes.

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/app/alphafold/run_alphafold.py", line 422, in
app.run(main)
File "/opt/alphafoldenv/lib/python3.8/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/opt/alphafoldenv/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "/app/alphafold/run_alphafold.py", line 398, in main
predict_structure(
File "/app/alphafold/run_alphafold.py", line 198, in predict_structure
prediction_result = model_runner.predict(processed_feature_dict,
File "/app/alphafold/alphafold/model/model.py", line 167, in predict
result = self.apply(self.params, jax.random.PRNGKey(random_seed), feat)
File "/opt/alphafoldenv/lib/python3.8/site-packages/jax/interpreters/xla.py", line 893, in _execute_compiled
out_bufs = compiled.execute(input_bufs)
RuntimeError: Resource exhausted: Out of memory while trying to allocate 64793134928 bytes.