bugfix for latest jax, clear_mem() was giving error

sokrypton · Oct 9, 2022 · 2cb66db · 2cb66db
1 parent 694c00b
commit 2cb66db
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/colabdesign/shared/utils.py b/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])