diff --git a/scripts/setfit/run_fewshot_multilabel.py b/scripts/setfit/run_fewshot_multilabel.py
index 6438c923..d0e8ad90 100644
--- a/scripts/setfit/run_fewshot_multilabel.py
+++ b/scripts/setfit/run_fewshot_multilabel.py
@@ -100,6 +100,7 @@ def __init__(self, args: argparse.Namespace) -> None:
             self.args.model, max_seq_length=args.max_seq_length, add_normalization_layer=args.add_normalization_layer
         )
         self.model = self.model_wrapper.model
+        self.model_original_state = copy.deepcopy(self.model.state_dict())
 
     def get_classifier(self, sbert_model: SentenceTransformer) -> SKLearnWrapper:
         if self.args.classifier == "logistic_regression":
@@ -113,7 +114,7 @@ def get_classifier(self, sbert_model: SentenceTransformer) -> SKLearnWrapper:
 
     def train(self, data: Dataset) -> SKLearnWrapper:
         "Trains a SetFit model on the given few-shot training data."
-        self.model.load_state_dict(copy.deepcopy(self.model_wrapper.model_original_state))
+        self.model.load_state_dict(copy.deepcopy(self.model_original_state))
 
         x_train = data["text"]
         y_train = data.remove_columns("text").to_pandas().values
diff --git a/src/setfit/modeling.py b/src/setfit/modeling.py
index 10f37cec..d8c8fb33 100644
--- a/src/setfit/modeling.py
+++ b/src/setfit/modeling.py
@@ -1,4 +1,3 @@
-import copy
 import os
 from dataclasses import dataclass
 from pathlib import Path
@@ -41,7 +40,6 @@
 class SetFitBaseModel:
     def __init__(self, model, max_seq_length: int, add_normalization_layer: bool) -> None:
         self.model = SentenceTransformer(model)
-        self.model_original_state = copy.deepcopy(self.model.state_dict())
         self.model.max_seq_length = max_seq_length
 
         if add_normalization_layer:
@@ -208,7 +206,6 @@ def __init__(
         self.multi_target_strategy = multi_target_strategy
         self.l2_weight = l2_weight
 
-        self.model_original_state = copy.deepcopy(self.model_body.state_dict())
         self.normalize_embeddings = normalize_embeddings
 
     def fit(