Pytorch Geometric - Loss always 'NaN'


I am training a simple GCN model. For some reason, my loss is always NaN from the beginning.
Please let me know if you need the whole code and data.

Thank you!

Here is my data:
Num_graphs = 800
X (num_nodes, num_fea) = [64,1]
Y (num_nodes, num_fea) = [64,1]

Here is my model:

class Graph_GCN(nn.Module):

def __init__(self, nfeat, nhid, nclass):

    super(Graph_GCN, self).__init__()
    self.conv1 = GCNConv(nfeat, nhid)
    self.conv2 = GCNConv(nhid, nclass)

def forward(self, data):

    input = data.x.float()
    edge_index = data.edge_index
    X = self.conv1(input,edge_index)
    X = self.conv2(X, edge_index)

    return X