improve this model: import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import RGCNConv, FastRGCNConv

class HeteroVGAE(nn.Module):
    def __init__(self, embedder, embed_dim, hidden_channels, out_channels, num_edge_types, full_edge_index, full_edge_index_ind, full_edge_types, full_edge_types_ind, id2entity_train, id2entity_test):
        super(HeteroVGAE, self).__init__()
        self.embedder = embedder

        # RGCN layers with added dropout and batch normalization
        self.rgcn1 = FastRGCNConv(embed_dim, hidden_channels, num_relations=num_edge_types)
        self.bn1 = nn.BatchNorm1d(hidden_channels)
        self.rgcn2 = FastRGCNConv(hidden_channels, out_channels, num_relations=num_edge_types)
        self.bn2 = nn.BatchNorm1d(out_channels)
        self.dropout = nn.Dropout(0.5)  # Dropout for regularization

        # Variational layers
        self.fc_mu = nn.Linear(embed_dim, embed_dim)
        self.fc_logstd = nn.Linear(embed_dim, embed_dim)
        self.fc_mu_hidden = nn.Linear(embed_dim, embed_dim)  # Additional hidden layer

        # Edge-type embeddings
        self.edge_type_embeddings = nn.Embedding(num_edge_types, out_channels)



        self.full_edge_index = full_edge_index
        self.full_edge_types = full_edge_types
        self.full_edge_index_ind = full_edge_index_ind
        self.full_edge_types_ind = full_edge_types_ind

        self.id2entity_train = id2entity_train
        self.id2entity_test = id2entity_test

    def encode(self, node_indices, phase='train'):
        if phase == 'train':
            node_embeddings = self.embedder(node_indices, phase="train")
        else:
            node_embeddings = self.embedder(node_indices, phase="test")

        # Variational layers
        x = F.relu(self.fc_mu_hidden(node_embeddings))
        x = self.dropout(x)

        mu = self.fc_mu(x)
        logstd = self.fc_logstd(x)
        z = self.reparameterize(mu, logstd)

        return z, mu, logstd

    def reparameterize(self, mu, logstd):
        std = torch.exp(0.5 * logstd)
        eps = torch.randn_like(std)
        return mu + eps * std


    def decode(self, z, edge_index, edge_types):
        row, col = edge_index
        edge_type_embeddings = self.edge_type_embeddings(edge_types)
        return torch.sigmoid((z[row] * z[col] * edge_type_embeddings).sum(dim=-1))


    def forward(self, edge_index, edge_types, phase='train', use_mu=False):
        node_indices = torch.unique(edge_index)

        z, mu, logstd = self.encode(node_indices, phase=phase)
        if use_mu:
            z = mu

        if phase == "train":
            full_x = self.embedder(torch.arange(self.embedder.train_embeddings.shape[0]), phase="train")
        else:
            full_x = self.embedder(torch.arange(self.embedder.test_embeddings.shape[0]), phase="test")

        full_x_new = torch.zeros_like(full_x)
        full_x_new[node_indices] += z

        mask = torch.ones(full_x.size(0), dtype=torch.bool)
        mask[node_indices] = False
        full_x_new[mask] += full_x[mask]

        if phase == "train":
            full_edge_index = self.full_edge_index
            full_edge_types = self.full_edge_types
        else:
            full_edge_index = self.full_edge_index_ind
            full_edge_types = self.full_edge_types_ind

        x = F.relu(self.rgcn1(full_x_new, full_edge_index, full_edge_types))
        x = self.bn1(x)
        x = F.relu(self.rgcn2(x, full_edge_index, full_edge_types))
        x = self.bn2(x)

        return self.decode(x, edge_index, edge_types), mu, logstd