Merge pull request #11 from ghanvert/CPU-execution

Added support for CPU execution

Author: ghanvert

src/accmt/__init__.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from datetime import timedelta
 
 import torch
@@ -37,8 +38,14 @@ def allow_tf32(flag=True):
 
 allow_tf32()
 
+IS_CPU = bool(int(os.environ.get("ACCMT_CPU", 0)))
+IS_GPU = not IS_CPU
+
 _init_kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=86400))
 _dataloader_config = DataLoaderConfiguration(use_seedable_sampler=True)
 accelerator = Accelerator(
-    kwargs_handlers=[_init_kwargs], dataloader_config=_dataloader_config, step_scheduler_with_optimizer=False
+    kwargs_handlers=[_init_kwargs],
+    dataloader_config=_dataloader_config,
+    step_scheduler_with_optimizer=False,
+    cpu=IS_CPU,
 )