Hello. I am very new in PyTorch graph Neural Network. I have multiple graphs and I want to do node classification. I divided the graphs to 80% of them for training and 10% for evaluation and 10% for testing. I tried to use GCN to classify nodes. To do that, I implemented the following code:

class GCN(torch.nn.Module):

def **init**(self, in_channels, hidden_channels, out_channels):

super().**init**()

self.conv1 = GCNConv(in_channels, hidden_channels)

self.conv2 = GCNConv(hidden_channels, out_channels)

# Define the linear layer

self.linear = torch.nn.Linear(out_channels, out_channels)

```
def forward(self, x,edge_index,edge_weight):
# x, edge_index = data.x, data.edge_index
x = self.conv1(x, edge_index,edge_weight)
x = F.relu(x)
x = F.dropout(x, training=self.training)
x = self.conv2(x, edge_index, edge_weight)
# Apply the ReLU activation function
x = F.relu(x)
# Apply the linear layer
x = self.linear(x)
return F.log_softmax(x, dim=1)
```

device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)

in_channels = gInfo.getNodeNumberOfFeatures() # Number of input features per node

hidden_channels = gInfo.getNodeNumberOfFeatures()+4 # Number of output features for the first convolutional layer

max_num_classes = 2 # Since the target output can be either 0 or 1

# Initialize the GCN model with the maximum number of classes

model = GCN(in_channels, hidden_channels, max_num_classes).to(device)

# Use Binary Cross Entropy Loss for scenarios where target has only one class

binary_loss_op = torch.nn.BCEWithLogitsLoss()

criterion = torch.nn.BCELoss()

optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

def train():

model.train()

```
total_loss = 0
for data in train_loader:
data = data.to(device)
data.requires_grad_(True)
optimizer.zero_grad()
output = model(data.x, data.edge_index, data.edge_weight)
# Calculate loss
loss = binary_loss_op(output, data.y.squeeze(dim=1))
total_loss += loss.item() * data.num_graphs
loss.backward()
optimizer.step()
return total_loss / len(train_loader.dataset)
```

Now, here for the loss part, the output has 2 dimensions and the data.y is only 0 or 1. I want to know how can I output to become 0 or 1 so to be able to compare with data.y?