Bark

import torch device = 'cuda:0' print(f"=> Set CUDA device = {device}") number_of_inputs = 2 input_size = 576 output_size = 10 class NeuralNetwork(torch.nn.Module): def __init__(self, input_size, output_size): super(NeuralNetwork, self).__init__() self.linear_relu_stack = torch.nn.Sequential( torch.nn.Linear(input_size, input_size), torch.nn.ReLU(), torch.nn.Linear(input_size, input_size), torch.nn.ReLU(), torch.nn.Linear(input_size, output_size), torch.nn.Softmax(dim=1) ) def forward(self, x): y = self.linear_relu_stack(x) print("=> Model input size", x.size(), "Model output size", y.size()) return y.argmax(1) model = NeuralNetwork(input_size, output_size) model.to(device) print(model) X = torch.rand(number_of_inputs, input_size, device=device) y = model(X) print(f"=> Predicted class: {y}")