commit 723f40d94e0d4497b294455e490fd1b1307fcd6e
author Summer Deng <summerdeng@fb.com> Fri Oct 26 20:46:26 2018 -0700
committer Facebook Github Bot <facebook-github-bot@users.noreply.github.com> Fri Oct 26 20:48:18 2018 -0700
tree 882ddd02c6661c36ff6988bb0abd1e219016e7a7
parent dae7616078414e162863de73ae3abf54930292d7

video model test workflow on CPU (#13203)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/13203

Minor changes in the test workflow to run the model on CPUs

Reviewed By: stephenyan1231

Differential Revision: D9925797

fbshipit-source-id: b7b1fb2658ab68b1ffc2b1f7b314958ea4732b32

caffe2/python/data_parallel_model.py

diff --git a/caffe2/python/data_parallel_model.py b/caffe2/python/data_parallel_model.py
index 954cf66..1ea110e 100644
--- a/caffe2/python/data_parallel_model.py
+++ b/caffe2/python/data_parallel_model.py
@@ -907,7 +907,7 @@
     return first_gpu_blobs.union(iteration_blobs)
 
 
-def FinalizeAfterCheckpoint(model, blobs=None):
+def FinalizeAfterCheckpoint(model, blobs=None, cpu_mode=False):
     '''
     This function should be called after loading parameters from a
     checkpoint / initial parameters file.
@@ -936,12 +936,14 @@
                 model._device_grouped_blobs[name] = grouped
 
         model._checkpoint_net = core.Net("checkpoint_sync_net")
-        model._checkpoint_net.RunAllOnGPU()
+        if not cpu_mode:
+            model._checkpoint_net.RunAllOnGPU()
 
         checkpoint_init_net = None
         if (model._rendezvous is not None and model._rendezvous['num_shards'] > 1):
             checkpoint_init_net = core.Net("checkpoint_init_net")
-            checkpoint_init_net.RunAllOnGPU()
+            if not cpu_mode:
+                checkpoint_init_net.RunAllOnGPU()
 
         _SyncAllParams(
             devices,