zpn commited on
Commit
86125d0
1 Parent(s): 5a2be72

fix: loading from pretrained

Browse files
Files changed (1) hide show
  1. modeling_hf_nomic_bert.py +4 -1
modeling_hf_nomic_bert.py CHANGED
@@ -330,7 +330,10 @@ class NomicBertPreTrainedModel(PreTrainedModel):
330
  if "add_pooling_layer" in kwargs:
331
  model = cls(config, *inputs, add_pooling_layer=kwargs.pop("add_pooling_layer"))
332
  else:
333
- model = cls(config, *inputs, add_pooling_layer=False)
 
 
 
334
  # TODO: fix this
335
  # Assuming we know what we're doing when loading from disk
336
  # Prob a bad assumption but i'm tired and want to train this asap
 
330
  if "add_pooling_layer" in kwargs:
331
  model = cls(config, *inputs, add_pooling_layer=kwargs.pop("add_pooling_layer"))
332
  else:
333
+ if cls == NomicBertModel:
334
+ model = cls(config, *inputs, add_pooling_layer=False)
335
+ else:
336
+ model = cls(config, *inputs)
337
  # TODO: fix this
338
  # Assuming we know what we're doing when loading from disk
339
  # Prob a bad assumption but i'm tired and want to train this asap