Fix Imports & Refactoring of test_gpu script

[dinosaurtirex]

Just removed outdated imports for a little bit better performance and refactored test_gpu.py script