diff --git a/build/lib/glycowork/ml/models.py b/build/lib/glycowork/ml/models.py
index 88802e5c..346112b0 100644
--- a/build/lib/glycowork/ml/models.py
+++ b/build/lib/glycowork/ml/models.py
@@ -291,14 +291,14 @@ def forward(self, prot, nodes, edge_index, batch, inference = False):
     # Fully connected part
     h_n = self.act1_n(self.bn1_n(self.fc1_n(h_n)))
 
-    x1 = self.fc2(self.dp1(h_n))
-    x2 = self.fc2(self.dp1(h_n))
-    x3 = self.fc2(self.dp1(h_n))
-    x4 = self.fc2(self.dp1(h_n))
-    x5 = self.fc2(self.dp1(h_n))
-    x6 = self.fc2(self.dp1(h_n))
-    x7 = self.fc2(self.dp1(h_n))
-    x8 = self.fc2(self.dp1(h_n))
+    x1 = self.fc2_n(self.dp1(h_n))
+    x2 = self.fc2_n(self.dp1(h_n))
+    x3 = self.fc2_n(self.dp1(h_n))
+    x4 = self.fc2_n(self.dp1(h_n))
+    x5 = self.fc2_n(self.dp1(h_n))
+    x6 = self.fc2_n(self.dp1(h_n))
+    x7 = self.fc2_n(self.dp1(h_n))
+    x8 = self.fc2_n(self.dp1(h_n))
 
     out = self.sigmoid(torch.mean(torch.stack([x1, x2, x3, x4, x5, x6, x7, x8]), dim = 0))
 
@@ -345,25 +345,25 @@ def prep_model(model_type, num_classes, libr = None,
         model = SweetNet(len(libr), num_classes = num_classes)
         model = model.apply(lambda module: init_weights(module, mode = 'sparse'))
         if trained:
-            model.load_state_dict(torch.load(trained_SweetNet))
+            model.load_state_dict(torch.load(trained_SweetNet, map_location = device))
         model = model.to(device)
     elif model_type == 'LectinOracle':
         model = LectinOracle(len(libr), num_classes = num_classes)
         model = model.apply(lambda module: init_weights(module, mode = 'xavier'))
         if trained:
-            model.load_state_dict(torch.load(trained_LectinOracle))
+            model.load_state_dict(torch.load(trained_LectinOracle, map_location = device))
         model = model.to(device)
     elif model_type == 'LectinOracle_flex':
         model = LectinOracle_flex(len(libr), num_classes = num_classes)
         model = model.apply(lambda module: init_weights(module, mode = 'xavier'))
         if trained:
-            model.load_state_dict(torch.load(trained_LectinOracle_flex))
+            model.load_state_dict(torch.load(trained_LectinOracle_flex, map_location = device))
         model = model.to(device)
     elif model_type == 'NSequonPred':
         model = NSequonPred()
         model = model.apply(lambda module: init_weights(module, mode = 'xavier'))
         if trained:
-            model.load_state_dict(torch.load(trained_NSequonPred))
+            model.load_state_dict(torch.load(trained_NSequonPred, map_location = device))
         model = model.to(device)
     else:
         print("Invalid Model Type")
diff --git a/build/lib/glycowork/motif/tokenization.py b/build/lib/glycowork/motif/tokenization.py
index bdd6cc63..c1684014 100644
--- a/build/lib/glycowork/motif/tokenization.py
+++ b/build/lib/glycowork/motif/tokenization.py
@@ -14,8 +14,9 @@
 from glycowork.motif.processing import min_process_glycans, canonicalize_iupac
 from glycowork.motif.graph import compare_glycans, glycan_to_nxGraph, graph_to_string
 
-chars = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T',
-         'V', 'W', 'Y', 'X', 'Z'] + ['z']
+chars = {'A':1, 'B':2, 'C':3, 'D':4, 'E':5, 'F':6, 'G':7, 'H':8, 'I':9, 'J':10, 'K':11,
+         'L':12, 'M':13, 'N':14, 'P':15, 'Q':16, 'R':17, 'S':18, 'T':19,
+         'V':20, 'W':21, 'Y':22, 'X':23, 'Z':24, 'z':25}
 
 io = pkg_resources.resource_stream(__name__, "mz_to_composition.csv")
 mapping_file = pd.read_csv(io)
@@ -35,7 +36,7 @@ def constrain_prot(proteins, libr = None):
   if libr is None:
     libr = chars
   # Check whether any character is not in libr and replace it with a 'z' placeholder character
-  forbidden = [k for k in set(list(''.join(proteins))) if k not in libr]
+  forbidden = [k for k in set(list(''.join(proteins))) if k not in libr.keys()]
   for k in forbidden:
     proteins = [j.replace(k, 'z') for j in proteins]
   return proteins
diff --git a/glycowork.egg-info/PKG-INFO b/glycowork.egg-info/PKG-INFO
index fca0387a..3e02ecd2 100644
--- a/glycowork.egg-info/PKG-INFO
+++ b/glycowork.egg-info/PKG-INFO
@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: glycowork
-Version: 0.8.0
+Version: 0.8.1
 Summary: Package for processing and analyzing glycans
 Home-page: https://github.com/BojarLab/glycowork
 Author: Daniel Bojar
@@ -70,10 +70,10 @@ alternative: <br>
 Note that we have optional extra installs for specialized use (even
 further instructions can be found in the `Examples` tab), such as: <br>
 *deep learning* <br> `pip install glycowork[ml]` <br> *drawing glycan
-images with GlycoDraw* <br> `pip install glycowork[draw]` <br>
-*analyzing atomic/chemical properties of glycans* <br>
-`pip install glycowork[chem]` <br> *everything* <br>
-`pip install glycowork[all]` <br>
+images with GlycoDraw (see install instructions in the `Examples` tab)*
+<br> `pip install glycowork[draw]` <br> *analyzing atomic/chemical
+properties of glycans* <br> `pip install glycowork[chem]` <br>
+*everything* <br> `pip install glycowork[all]` <br>
 
 ## Data & Models
 
diff --git a/glycowork/ml/models.py b/glycowork/ml/models.py
index 88802e5c..346112b0 100644
--- a/glycowork/ml/models.py
+++ b/glycowork/ml/models.py
@@ -291,14 +291,14 @@ def forward(self, prot, nodes, edge_index, batch, inference = False):
     # Fully connected part
     h_n = self.act1_n(self.bn1_n(self.fc1_n(h_n)))
 
-    x1 = self.fc2(self.dp1(h_n))
-    x2 = self.fc2(self.dp1(h_n))
-    x3 = self.fc2(self.dp1(h_n))
-    x4 = self.fc2(self.dp1(h_n))
-    x5 = self.fc2(self.dp1(h_n))
-    x6 = self.fc2(self.dp1(h_n))
-    x7 = self.fc2(self.dp1(h_n))
-    x8 = self.fc2(self.dp1(h_n))
+    x1 = self.fc2_n(self.dp1(h_n))
+    x2 = self.fc2_n(self.dp1(h_n))
+    x3 = self.fc2_n(self.dp1(h_n))
+    x4 = self.fc2_n(self.dp1(h_n))
+    x5 = self.fc2_n(self.dp1(h_n))
+    x6 = self.fc2_n(self.dp1(h_n))
+    x7 = self.fc2_n(self.dp1(h_n))
+    x8 = self.fc2_n(self.dp1(h_n))
 
     out = self.sigmoid(torch.mean(torch.stack([x1, x2, x3, x4, x5, x6, x7, x8]), dim = 0))
 
@@ -345,25 +345,25 @@ def prep_model(model_type, num_classes, libr = None,
         model = SweetNet(len(libr), num_classes = num_classes)
         model = model.apply(lambda module: init_weights(module, mode = 'sparse'))
         if trained:
-            model.load_state_dict(torch.load(trained_SweetNet))
+            model.load_state_dict(torch.load(trained_SweetNet, map_location = device))
         model = model.to(device)
     elif model_type == 'LectinOracle':
         model = LectinOracle(len(libr), num_classes = num_classes)
         model = model.apply(lambda module: init_weights(module, mode = 'xavier'))
         if trained:
-            model.load_state_dict(torch.load(trained_LectinOracle))
+            model.load_state_dict(torch.load(trained_LectinOracle, map_location = device))
         model = model.to(device)
     elif model_type == 'LectinOracle_flex':
         model = LectinOracle_flex(len(libr), num_classes = num_classes)
         model = model.apply(lambda module: init_weights(module, mode = 'xavier'))
         if trained:
-            model.load_state_dict(torch.load(trained_LectinOracle_flex))
+            model.load_state_dict(torch.load(trained_LectinOracle_flex, map_location = device))
         model = model.to(device)
     elif model_type == 'NSequonPred':
         model = NSequonPred()
         model = model.apply(lambda module: init_weights(module, mode = 'xavier'))
         if trained:
-            model.load_state_dict(torch.load(trained_NSequonPred))
+            model.load_state_dict(torch.load(trained_NSequonPred, map_location = device))
         model = model.to(device)
     else:
         print("Invalid Model Type")
diff --git a/glycowork/motif/tokenization.py b/glycowork/motif/tokenization.py
index bdd6cc63..c1684014 100644
--- a/glycowork/motif/tokenization.py
+++ b/glycowork/motif/tokenization.py
@@ -14,8 +14,9 @@
 from glycowork.motif.processing import min_process_glycans, canonicalize_iupac
 from glycowork.motif.graph import compare_glycans, glycan_to_nxGraph, graph_to_string
 
-chars = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T',
-         'V', 'W', 'Y', 'X', 'Z'] + ['z']
+chars = {'A':1, 'B':2, 'C':3, 'D':4, 'E':5, 'F':6, 'G':7, 'H':8, 'I':9, 'J':10, 'K':11,
+         'L':12, 'M':13, 'N':14, 'P':15, 'Q':16, 'R':17, 'S':18, 'T':19,
+         'V':20, 'W':21, 'Y':22, 'X':23, 'Z':24, 'z':25}
 
 io = pkg_resources.resource_stream(__name__, "mz_to_composition.csv")
 mapping_file = pd.read_csv(io)
@@ -35,7 +36,7 @@ def constrain_prot(proteins, libr = None):
   if libr is None:
     libr = chars
   # Check whether any character is not in libr and replace it with a 'z' placeholder character
-  forbidden = [k for k in set(list(''.join(proteins))) if k not in libr]
+  forbidden = [k for k in set(list(''.join(proteins))) if k not in libr.keys()]
   for k in forbidden:
     proteins = [j.replace(k, 'z') for j in proteins]
   return proteins
diff --git a/setup.py b/setup.py
index e15ea452..e3c624c2 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="glycowork",
-    version="0.8.0",
+    version="0.8.1",
     author="Daniel Bojar",
     author_email="daniel.bojar@gu.se",
     description="Package for processing and analyzing glycans",