lausannel commited on
Commit
e50cd48
1 Parent(s): d9ab9eb

fix: change position of remove offload path and fix params in __init__()

Browse files
Files changed (1) hide show
  1. src/backend/moe_infinity.py +5 -3
src/backend/moe_infinity.py CHANGED
@@ -31,15 +31,17 @@ class MoEHFLM(HFLMWithMeasurement):
31
  self.use_chat_template = use_chat_template
32
  if "device" in kwargs:
33
  kwargs.pop("device")
 
 
34
  super().__init__(
35
- *args, **kwargs, pretrained=pretrained, device_map="cuda:0"
36
  ) # Assuming HFLM accepts a 'pretrained' arg and handles it
37
  # self._create_model()
38
- shutil.rmtree(os.path.join(self.offload_path, "moe-infinity-offloads"))
39
 
40
  def __del__(self):
41
  # Clean up offloaded models from self.offload_path
42
- shutil.rmtree(os.path.join(self.offload_path, "moe-infinity-offloads"))
 
43
 
44
  def _create_model(self, *args, **kwargs):
45
  """
 
31
  self.use_chat_template = use_chat_template
32
  if "device" in kwargs:
33
  kwargs.pop("device")
34
+ if os.path.exists(os.path.join(self.offload_path, "moe-infinity-offloads")):
35
+ shutil.rmtree(os.path.join(self.offload_path, "moe-infinity-offloads"))
36
  super().__init__(
37
+ *args, **kwargs, pretrained=pretrained, device="cuda:0"
38
  ) # Assuming HFLM accepts a 'pretrained' arg and handles it
39
  # self._create_model()
 
40
 
41
  def __del__(self):
42
  # Clean up offloaded models from self.offload_path
43
+ if os.path.exists(os.path.join(self.offload_path, "moe-infinity-offloads")):
44
+ shutil.rmtree(os.path.join(self.offload_path, "moe-infinity-offloads"))
45
 
46
  def _create_model(self, *args, **kwargs):
47
  """