Thanks.
I’ve just faced a problem when applied your code to my work.
I met the error like ‘RuntimeError: expected device cuda:0 and dtype Float but got device cpu and dtype Float’ at line 223.
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    mask = torch.tensor([[1, 0, 0], [1, 1, 0], [1, 0, 1]], dtype=torch.float)
    model = neuralNet(input_size, hidden_size, num_classes, mask)
    model.to(device)
    ...
    # Move tensors to the configured device
    inputs = inputs.to(device)
    labels = labels.to(device)

    # Forward pass
    y_pred = model(inputs)
    loss = criterion(y_pred, labels)


    optimizer.zero_grad()
    loss.backward()

    # Zero out gradients
    with torch.no_grad():
        model.fc1.weight.grad.mul_(mask) ---> line 223
        model.fc2.weight.grad.mul_(mask)
        model.fc3.weight.grad.mul_(mask)
        model.fc4.weight.grad.mul_(mask)

    optimizer.step()

How can I solve this problem?