fix: do not always set classification to false

sparks-baird · Feb 9, 2023 · eb974a1 · eb974a1
1 parent 223dd7a
commit eb974a1
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/crabnet/crabnet_.py b/crabnet/crabnet_.py
@@ -245,14 +245,13 @@ def __init__(
         self.max_lr = max_lr
 
         # Apply BCEWithLogitsLoss to model output if binary classification is True
-        if classification:
-            self.classification = True
+        self.classification = classification
+
 
         self.model_name = model_name
         self.mat_prop = mat_prop
         self.data_loader = None
         self.train_loader = None
-        self.classification = False
         self.n_elements = n_elements
 
         self.fudge = fudge  #  expected fractional tolerance (std. dev) ~= 2%