no need for large number of steps in test

ml-agents/mlagents/trainers/tests/torch/test_attention.py

         weight_decay=1e-6,
     )
     batch_size = 20
-    for _ in range(5000):
+    for _ in range(5):
         center = torch.rand((batch_size, size))
         key = torch.rand((batch_size, n_k, size))
         with torch.no_grad():