Add missing Pyre mode headers] [batch:1/N] [A] [A] (meta-pytorch#1016)

Summary:
X-link: meta-pytorch/torchrec#3204

X-link: meta-pytorch/torchx#1091


Differential Revision: D78138183

Author: Pyre Bot Jr.

tests/utils/test_device_mesh.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 import unittest
 
 from torchtnt.utils.device_mesh import (

torchtnt/utils/device_mesh.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 from typing import Optional
 
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh