Updated profiler examples to be nonzcc

aws · Sep 7, 2022 · bb64f33 · bb64f33
1 parent fef2253
commit bb64f33
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 0 deletions.
diff --git a/sagemaker-debugger/pytorch_profiling/entry_point/pt_res50_cifar10_distributed.py b/sagemaker-debugger/pytorch_profiling/entry_point/pt_res50_cifar10_distributed.py
@@ -53,6 +53,7 @@ def train(batch_size, epoch, net, hook, device, local_rank):
     epoch_times = []
 
     if hook:
+        hook.register_module(net)
         hook.register_loss(loss_optim)
     # train the model
 

diff --git a/sagemaker-debugger/pytorch_profiling/entry_point/pt_res50_cifar10_horovod_dataloader.py b/sagemaker-debugger/pytorch_profiling/entry_point/pt_res50_cifar10_horovod_dataloader.py
@@ -101,6 +101,7 @@ def train(batch_size, epoch, net, hook, args, local_rank):
 
         print("START VALIDATING")
         if hook:
+            hook.register_module(net)
             hook.set_mode(modes.EVAL)
         test_sampler.set_epoch(i)
         net.eval()

diff --git a/sagemaker-debugger/pytorch_profiling/entry_point/pytorch_res50_cifar10_dataloader.py b/sagemaker-debugger/pytorch_profiling/entry_point/pytorch_res50_cifar10_dataloader.py
@@ -75,6 +75,7 @@ def train(args, net, device):
     epoch_times = []
 
     if hook:
+        hook.register_module(net)
         hook.register_loss(loss_optim)
     # train the model