Fix device issues

custom_interface.py

@@ -80,7 +80,8 @@ class Discrete_EmbeddingLayer(torch.nn.Module):
         self.init = init
         self.layers = layers
         self.available_layers = available_layers
-        self.offsets = self.build_offsets()
+        self.register_buffer("offsets", self.build_offsets())
+        self.register_buffer("layer_embs", self.compute_layer_embs())
         self.chunk_size = chunk_size
 
     def init_embedding(self, weights):