Problem with step four in colabpro

I keep getting the following error message in step 4:
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
[](https://7c9p1ldk636-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240408-060136_RC00_622797729#) in ()
25 import py3Dmol
26
---> 27 from alphafold.model import model
28 from alphafold.model import config
29 from alphafold.model import data

7 frames
[/usr/local/lib/python3.10/dist-packages/jax/_src/deprecations.py](https://7c9p1ldk636-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20240408-060136_RC00_622797729#) in getattr(name)
52 warnings.warn(message, DeprecationWarning, stacklevel=2)
53 return fn
---> 54 raise AttributeError(f"module {module!r} has no attribute {name!r}")
55
56 return getattr

AttributeError: module 'jax' has no attribute 'linear_util'

I have added in the suggested code but it doesn't resolve the problem. Please advise