diff --git a/plugins/train/model/dfl_h128.py b/plugins/train/model/dfl_h128.py
index 3afcf408c4..a11ef6d983 100644
--- a/plugins/train/model/dfl_h128.py
+++ b/plugins/train/model/dfl_h128.py
@@ -27,7 +27,7 @@ def encoder(self):
         var_x = Dense(8 * 8 * self.encoder_dim)(var_x)
         var_x = Reshape((8, 8, self.encoder_dim))(var_x)
         var_x = UpscaleBlock(self.encoder_dim)(var_x)
-        return KerasModel(input_, var_x)
+        return KerasModel(input_, var_x, name=self.name)
 
     def decoder(self, side):
         """ DFL H128 Decoder """