Getting model logits is very slow

[henrikfo]

Issue

Problem Description

After each shadow model is done training, it logits of the dataset is calculated and fetched.
This process takes approximately 6-7x the time one whole epoch takes to train.
The fault the model and logits are swapping devices multiple times.

Expected Behavior

It should be faster than training one epoch, since it is just doing inference.

What Needs to be Done

change ModelLogits in signal.py and PytorchModel.get_logits() in signal_extractor.py

How Can It Be Tested or Reproduced

Check the time it takes to train a shadow model one epoch then watch the time it takes to fetch the logits