for nb_hidden in [16, 32, 64, 128, 256, 512, 1024]:
for linear_layer in [nn.Linear, QLinear]:
+ # The model
+
model = nn.Sequential(
nn.Flatten(),
linear_layer(784, nb_hidden),
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
- ######################################################################
+ #
for k in range(nb_epochs):
- ############################################
# Train
model.train()
loss.backward()
optimizer.step()
- ############################################
# Test
model.eval()