Getting model logits is very slow

# 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