Problem with step four in colabpro

[ScottStro]

I keep getting the following error message in step 4:

AttributeError Traceback (most recent call last)
in <cell line: 27>()
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 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