Enable QAT for static llama definition (pytorch#13285)

Summary:

The model needed small modifications to be able to run QAT on GPUs.

Reviewed By: YIWENX14

Differential Revision: D79841467

Author: navsud

examples/models/llama/rope.py

@@ -47,9 +47,10 @@ def precompute_freqs_cis(
     use_scaled: bool = False,
     scale_factor: Optional[int] = None,
     high_freq_factor: int = 4,
+    device: torch.device = torch.device("cpu"),
 ):
     freqs = 1.0 / (
-        theta ** (torch.arange(0, dim, 2, device="cpu")[: (dim // 2)].float() / dim)
+        theta ** (torch.arange(0, dim, 2, device=device)[: (dim // 2)].float() / dim)
     )
     t = torch.arange(end, device=freqs.device)  # pyre-ignore
     if use_scaled: