import numpy as np
import torch
import torch.nn.functional as F
import torch.nn as nn

from models import GCN
from datasets import DDInteractionDataset



if __name__ == '__main__':
    ddiDataset = DDInteractionDataset
    model = GCN(dataset.num_features, dataset.num_classes)
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    # training on CPU
    for epoch in range(1, 6):
        optimizer.zero_grad()
        out = model(data.x, data.edge_index, data.edge_attr)
        loss = F.cross_entropy(out, data.y)
        loss.backward()
        optimizer.step()
        print(f"Epoch: {epoch}, Loss: {loss}")