Thank you for the answer, any suggestion on which optimizer & loss function to choose? does it really matter in this case? i am currently using :
    criterion = nn.BCELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)