Output of custom linear model for Hetero Data in GraphGym is not as expected

[mdanb]

I registered the following Linear module in GraphGym:

from torch_geometric.nn import Linear as Linear_pyg

@register_layer('custom_linear')
class Linear(nn.Module):
    def __init__(self, layer_config: LayerConfig):
        super().__init__()
        self.model = Linear_pyg(layer_config.dim_in, layer_config.dim_out,
                                bias=layer_config.has_bias)

    def forward(self, batch):
        if isinstance(batch, torch.Tensor):
            batch = self.model(batch)
        else:
            batch.x_dict["node"] = self.model(batch.x_dict["node"])
        return batch

layer_config.dim_in is 200 and layer_config.dim_out is 256. However, when I step through the code to check batch.x_dict["node"].shape after the line batch.x_dict["node"] = self.model(batch.x_dict["node"]), I see that the dimensionality of the data is still 200, not 256. Interestingly, when I check self.model(batch.x_dict["node"]), it is 256 dimensional. But the line batch.x_dict["node"] = self.model(batch.x_dict["node"]) does not change batch.x_dict["node"]

PS: batch["node"].x = self.model(batch.x_dict["node"]) works on the other hand.
PPS: I'm using HeteroData because I have multiple edge types (but only one node type, which is why I can do batch.x_dict["node"] = self.model(batch.x_dict["node"]))

[akihironitta]

Just to make sure I understand, it'd be nice if you could provide a complete script with print/assert statements to show what the correct/wrong shape of the output is.