@@ -238,19 +238,18 @@ def mask_attr_prediction(self, g, x):
         else:
             use_g = pre_use_g
 
-        # use_g = pre_use_g
         enc_rep, all_hidden = self.encoder(use_g, use_x, return_hidden=True)
         if self._concat_hidden:
             enc_rep = torch.cat(all_hidden, dim=1)
 
         # ---- attribute reconstruction ----
         rep = self.encoder_to_decoder(enc_rep)
 
-        if self._decoder_type != "mlp":
+        if self._decoder_type not in ("mlp", "linear"):
             # * remask, re-mask
-            rep[mask_nodes] = 0 #
+            rep[mask_nodes] = 0
 
-        if self._decoder_type == "mlp":
+        if self._decoder_type in ("mlp", "liear") :
             recon = self.decoder(rep)
         else:
             recon = self.decoder(pre_use_g, rep)