bugfix for latest jax, clear_mem() was giving error

Author: sokrypton

colabdesign/shared/utils.py

@@ -19,7 +19,10 @@ def clear_mem():
       for obj_name in dir(module):
         obj = getattr(module, obj_name)
         if hasattr(obj, "cache_clear"):
-          obj.cache_clear()
+          try:
+            obj.cache_clear()
+          except:
+            pass
   gc.collect()
 
 def update_dict(D, *args, **kwargs):
@@ -116,4 +119,4 @@ def copy_missing(a,b):
     if i not in b:
       b[i] = v
     elif isinstance(v,dict):
-      copy_missing(v,b[i])
+      copy_missing(v,b[i])