diff --git a/pytorch/resnet/image_classification/dataloaders.py b/pytorch/resnet/image_classification/dataloaders.py
index 879df44..5fdd4bf 100644
--- a/pytorch/resnet/image_classification/dataloaders.py
+++ b/pytorch/resnet/image_classification/dataloaders.py
@@ -34,7 +34,8 @@
 import torchvision.transforms as transforms
 from PIL import Image
 from functools import partial
-import smdistributed.dataparallel.torch.distributed as dist
+#import smdistributed.dataparallel.torch.distributed as dist
+import torch.distributed as dist
 
 from image_classification.autoaugment import AutoaugmentImageNetPolicy
 
diff --git a/pytorch/resnet/image_classification/training.py b/pytorch/resnet/image_classification/training.py
index 6f07ea2..5184752 100644
--- a/pytorch/resnet/image_classification/training.py
+++ b/pytorch/resnet/image_classification/training.py
@@ -41,8 +41,13 @@
 
 from .optimizers import get_sgd_optimizer, get_rmsprop_optimizer
 from .models.common import EMA
-import smdistributed.dataparallel.torch.distributed as dist
-from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP
+#import smdistributed.dataparallel.torch.distributed as dist
+#from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch_smddp
+dist.init_process_group(backend='smddp')
+
 from torch.cuda.amp import autocast
 
 ACC_METADATA = {"unit": "%", "format": ":.2f"}
@@ -480,7 +485,7 @@ def train_loop(
         epochs_since_improvement = 0
     backup_prefix = checkpoint_filename[:-len("checkpoint.pth.tar")] if \
         checkpoint_filename.endswith("checkpoint.pth.tar") else ""
-    
+
     print(f"RUNNING EPOCHS FROM {start_epoch} TO {end_epoch}")
     with utils.TimeoutHandler() as timeout_handler:
         interrupted = False
diff --git a/pytorch/resnet/image_classification/utils.py b/pytorch/resnet/image_classification/utils.py
index 548e1e4..bfb9fe6 100644
--- a/pytorch/resnet/image_classification/utils.py
+++ b/pytorch/resnet/image_classification/utils.py
@@ -33,7 +33,8 @@
 import torch
 import shutil
 import signal
-import smdistributed.dataparallel.torch.distributed as dist
+#import smdistributed.dataparallel.torch.distributed as dist
+import torch.distributed as dist
 
 
 def should_backup_checkpoint(args):
diff --git a/pytorch/resnet/main.py b/pytorch/resnet/main.py
index f8cacc1..0c05e09 100644
--- a/pytorch/resnet/main.py
+++ b/pytorch/resnet/main.py
@@ -40,8 +40,9 @@
 import torchvision.transforms as transforms
 import torchvision.datasets as datasets
 
-import smdistributed.dataparallel.torch.distributed as dist
-dist.init_process_group()
+#import smdistributed.dataparallel.torch.distributed as dist
+#dist.init_process_group()
+import torch.distributed as dist
 
 import image_classification.logger as log
 
@@ -330,7 +331,7 @@ def prepare_for_training(args, model_args, model_arch):
     args.distributed = False
     if dist.get_world_size() > 1:
         args.distributed = True
-        args.local_rank = dist.get_local_rank()
+        args.local_rank = dist.get_rank() % 8
     else:
         args.local_rank = 0
 
@@ -338,7 +339,7 @@ def prepare_for_training(args, model_args, model_arch):
     args.world_size = 1
 
     if args.distributed:
-        args.gpu = dist.get_local_rank()
+        args.gpu = dist.get_rank() % 8
         torch.cuda.set_device(args.gpu)
         # dist.init_process_group(backend="nccl", init_method="env://")
         args.world_size = dist.get_world_size()
@@ -606,7 +607,7 @@ def main(args, model_args, model_arch):
     add_parser_arguments(parser)
 
     args, rest = parser.parse_known_args()
-    
+
     model_arch = available_models()[args.arch]
     model_args, rest = model_arch.parser().parse_known_args(rest)
     print(model_args)