Fix benchmarks

bigscience-workshop · Apr 10, 2024 · 16d97fc · 16d97fc
1 parent a447cbe
commit 16d97fc
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/src/petals/models/mixtral/model.py b/src/petals/models/mixtral/model.py
@@ -135,9 +135,7 @@ def ln_f(self) -> nn.Module:  # For compatibility with RemoteGenerationMixin in
         return self.norm
 
 
-class DistributedMixtralForCausalLM(
-    DefaultRevisionMixin, FromPretrainedMixin, RemoteGenerationMixin, MixtralForCausalLM
-):
+class DistributedMixtralForCausalLM(FromPretrainedMixin, RemoteGenerationMixin, MixtralForCausalLM):
     _keys_to_ignore_on_load_missing = DistributedMixtralModel._keys_to_ignore_on_load_missing
     _keys_to_ignore_on_load_unexpected = DistributedMixtralModel._keys_to_ignore_on_load_unexpected
 
@@ -159,9 +157,7 @@ def transformer(self) -> DistributedMixtralModel:  # For compatibility with Remo
         return self.model
 
 
-class DistributedMixtralForSequenceClassification(
-    DefaultRevisionMixin, FromPretrainedMixin, MixtralForSequenceClassification
-):
+class DistributedMixtralForSequenceClassification(FromPretrainedMixin, MixtralForSequenceClassification):
     def __init__(self, config: DistributedMixtralConfig):
         MixtralPreTrainedModel.__init__(self, config)
         self.num_labels = config.num_labels