The code looks alright regarding the GPU usage.
You could profile it and see where the bottleneck is. Given that your model is quite small, you might see the kernel launch overheads and could increase the workload by e.g. using a larger model or larger batch size. Also, you could check CUDA graphs to see if this could reduce the launch overheads.
Some issues in the code unrelated to the device usage:

don’t use the forward method explicitly as it will skip potentially registered hooks. Use output = model(input) instead.

loss_train[epoch] += loss will store the loss tensor with the attached computation graph, which could eventually yield an out of memory issue. .detach() the tensor before accumulating it or call .item() on it.
I would write the model in a device-agnostic way and avoid passing the device to the constructor. If you really want to keep the test method inside the model, use .to(next(self.parameters()).device)