Add import.

Author: Xreki

graph_net/torch/sample_passes/agent_unittest_generator.py

@@ -24,6 +24,7 @@
 {%- endif -%}
 {{"\n"}}
 import torch
+from torch import device
 
 
 {% macro get_input_tensor_instance(tensor_meta, device) -%}