diff --git a/autopeptideml/apml.py b/autopeptideml/apml.py
index b8f85c7..eb7dc89 100644
--- a/autopeptideml/apml.py
+++ b/autopeptideml/apml.py
@@ -24,7 +24,7 @@
 from .reps import RepEngineBase, PLMs, CLMs, FPs
 
 
-__version__ = '2.0.4'
+__version__ = '2.0.6'
 
 
 class AutoPeptideML:
diff --git a/autopeptideml/main.py b/autopeptideml/main.py
index a60bfaf..af1ee3e 100644
--- a/autopeptideml/main.py
+++ b/autopeptideml/main.py
@@ -239,17 +239,18 @@ def predict(result_dir: str, features_path: str,
     if n_jobs == -1:
         n_jobs = cpu_count()
 
-    if not osp.isfile(metadata_path):
-        raise RuntimeError("No metadata file was found. Check you're using the correct AutoPeptideML version.",
-                           "Try: pip install autopeptideml<=2.0.1",
-                           "Alternatively, check that the result_dir path is properly formatted.")
+    if osp.isfile(metadata_path):
+        metadata = yaml.safe_load(open(metadata_path))
+        print(f"Using model created on {metadata['start-time']} with AutoPeptideML v.{metadata['autopeptideml-version']}")
+
+        # raise RuntimeError("No metadata file was found. Check you're using the correct AutoPeptideML version.",
+        #                    "Try: pip install autopeptideml<=2.0.1",
+        #                    "Alternatively, check that the result_dir path is properly formatted.")
     if not osp.isdir(ensemble_path):
         raise RuntimeError("No ensemble path was found in result_dir.  Check you're using the correct AutoPeptideML version.",
                            "Try: pip install autopeptideml<=2.0.1",
                            "Alternatively, check that the result_dir path is properly formatted.")
 
-    metadata = yaml.safe_load(open(metadata_path))
-    print(f"Using model created on {metadata['start-time']} with AutoPeptideML v.{metadata['autopeptideml-version']}")
     ensemble = VotingEnsemble.load(ensemble_path)
     df = read_data(features_path)
     if feature_field is None:
@@ -281,9 +282,10 @@ def predict(result_dir: str, features_path: str,
                               verbose=True)
     x = {}
     for rep in ensemble.reps:
+        if rep in x:
+            continue
         if rep in PLMs:
             from .reps.lms import RepEngineLM
-
             repengine = RepEngineLM(rep)
             repengine.move_to_device(device)
             x[rep] = repengine.compute_reps(
@@ -292,7 +294,6 @@ def predict(result_dir: str, features_path: str,
 
         elif rep in CLMs:
             from .reps.lms import RepEngineLM
-
             repengine = RepEngineLM(rep)
             repengine.move_to_device(device)
             x[rep] = repengine.compute_reps(
@@ -310,7 +311,7 @@ def predict(result_dir: str, features_path: str,
 
     try:
         preds, std = ensemble.predict_proba(x)
-        preds = preds[:, 1]
+        # preds = preds
     except AttributeError:
         preds, std = ensemble.predict(x)
     if 'pipe-seq' in df:
diff --git a/autopeptideml/train/architectures.py b/autopeptideml/train/architectures.py
index 29cad8c..444b74a 100644
--- a/autopeptideml/train/architectures.py
+++ b/autopeptideml/train/architectures.py
@@ -43,8 +43,12 @@ def predict(self, x: np.ndarray):
         :return: Predicted values from the ONNX model.
         :rtype: np.ndarray
         """
-        input_dict = {"float_input": x.astype(np.float32)}
-        preds = self.session.run(None, input_dict)
+        try:
+            input_dict = {"float_input": x.astype(np.float32)}
+            preds = self.session.run(None, input_dict)
+        except ValueError:
+            input_dict = {"X": x.astype(np.float32)}
+            preds = self.session.run(None, input_dict)
         return preds[0]
 
     def predict_proba(self, x: np.ndarray):
@@ -62,8 +66,12 @@ def predict_proba(self, x: np.ndarray):
         :return: Array of predicted probabilities for the positive class (shape: [n_samples]).
         :rtype: np.ndarray
         """
-        input_dict = {"float_input": x.astype(np.float32)}
-        preds = self.session.run(None, input_dict)
+        try:
+            input_dict = {"float_input": x.astype(np.float32)}
+            preds = self.session.run(None, input_dict)
+        except ValueError:
+            input_dict = {"X": x.astype(np.float32)}
+            preds = self.session.run(None, input_dict)
         return np.array([i[1] for i in preds[1]])
 
 
@@ -230,6 +238,7 @@ def load(self, path) -> "VotingEnsemble":
                 raise RuntimeError(f"File: {filepath} in path: {path} is not ONNX model.")
             models.append(OnnxModel(filepath))
             reps.append(osp.basename(filepath).split("_")[1].split('.')[0])
+        reps = [r if r != 'class' else 'esm2-8m' for r in reps]
         return VotingEnsemble(models, reps)
 
 
diff --git a/setup.py b/setup.py
index 666e698..ecd9298 100644
--- a/setup.py
+++ b/setup.py
@@ -73,6 +73,6 @@ def get_files_in_dir(path: Path, base: Path) -> list:
     name='autopeptideml',
     packages=find_packages(exclude=['examples']),
     url='https://ibm.github.io/AutoPeptideML/',
-    version='2.0.5',
+    version='2.0.6',
     zip_safe=False,
 )