Add device to Attention obj

Author: joecummings

torchtext/prototype/models/t5/modules.py

@@ -74,6 +74,8 @@ def __init__(
         else:
             self.relative_attention_bias = None
 
+        self.device = device
+
     def forward(
         self,
         query: Tensor,