Update graph_decoder/transformer.py

graph_decoder/transformer.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 from .layers import Attention, MLP
 from .conditions import TimestepEmbedder, ConditionEmbedder
-from .diffusion_utils import PlaceHolder
+# from .diffusion_utils import PlaceHolder
 
 def modulate(x, shift, scale):
     return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
@@ -98,7 +98,8 @@ class Transformer(nn.Module):
 
         # X: B * N * dx, E: B * N * N * de
         X, E = self.output_layer(X, X_in, E_in, c, t, node_mask)
-        return PlaceHolder(X=X, E=E, y=None).mask(node_mask)
+        # return PlaceHolder(X=X, E=E, y=None).mask(node_mask)
+        return X, E
 
 class Block(nn.Module):
     def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):