Update classifier.py

Social-Evolution-and-Behavior · Apr 2, 2020 · bd6d01a · bd6d01a
1 parent 14cef5a
commit bd6d01a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/antrax/classifier.py b/antrax/classifier.py
@@ -151,7 +151,8 @@ def save(self, modelfile):
         # save classifier params
         with h5py.File(modelfile, 'a') as f:
             f.create_dataset('prmtrs', data=json.dumps(self.prmtrs))
-            f.create_dataset('classes', (len(self.classes), 1), 'S10', [np.string_(x) for x in self.classes])
+            if self.classes is not None:
+                f.create_dataset('classes', (len(self.classes), 1), 'S10', [np.string_(x) for x in self.classes])
 
     def prepare_images(self):