cbensimon HF Staff commited on
Commit
7e1b70d
·
1 Parent(s): 1017ac0

Make ZeroGPUCompiledModel actually serializable

Browse files
Files changed (1) hide show
  1. optimization_utils.py +2 -0
optimization_utils.py CHANGED
@@ -45,6 +45,8 @@ class ZeroGPUCompiledModel:
45
  compiled_model.load_constants(self.weights.constants_map, check_full_update=True, user_managed=True)
46
  self.compiled_model.set(compiled_model)
47
  return compiled_model(*args, **kwargs)
 
 
48
 
49
 
50
  def aoti_compile(
 
45
  compiled_model.load_constants(self.weights.constants_map, check_full_update=True, user_managed=True)
46
  self.compiled_model.set(compiled_model)
47
  return compiled_model(*args, **kwargs)
48
+ def __reduce__(self):
49
+ return ZeroGPUCompiledModel, (self.archive_file, self.weights)
50
 
51
 
52
  def aoti_compile(