[Fix] Missed import and dim reference

Signed-off-by: Gene Der Su <[email protected]>

Author: GeneDer

setup.py

@@ -7,6 +7,7 @@
 """Installation script."""
 
 import os
+import sys
 import time
 from pathlib import Path
 from typing import List, Tuple

transformer_engine/jax/flax/transformer.py

@@ -637,7 +637,7 @@ def __call__(
             dropout_rng = self.make_rng(self.dropout_rng_name)
 
         if self.scale_factor is None:
-            scale_factor = 1.0 / sqrt(self.head_dim_qk)
+            scale_factor = 1.0 / sqrt(head_dim_qk)
         else:
             scale_factor = self.scale_factor
         del self.scale_factor