nnGraph multi-GPU Torch - deep-learning

Question

This question is about making any nnGraph network run on multiple GPUs and not specific to the following network instance
I am trying to train a network which is constructed with nnGraph. The backward diagram is attached. I am trying to run the parallelModel (see code or fig Node 9) in a multi-GPU setting. If I attach the parallel model to a nn.Sequential container and then create a DataParallelTable it works in a multi-GPU setting (without nnGraph). However, after attaching it to nnGraph I get an error. The backward pass works if I train on a single GPU (setting true to false in the if statements), but in a multi-GPU setting I get an error "gmodule.lua:418: attempt to index local 'gradInput' (a nil value)". I think Node 9 in backward pass should run on multiple-GPUs, however that's not happening. Creating DataParallelTable on nnGraph didn't work for me, however I thought atleast putting internal Sequential networks in a DataParallelTable would work. Is there some other way to split the initial data which is being passed to nnGraph so that it runs on multiple-GPUs?
require 'torch'
require 'nn'
require 'cudnn'
require 'cunn'
require 'cutorch'
require 'nngraph'
data1 = torch.ones(4,20):cuda()
data2 = torch.ones(4,10):cuda()
tmodel = nn.Sequential()
tmodel:add(nn.Linear(20,10))
tmodel:add(nn.Linear(10,10))
parallelModel = nn.ParallelTable()
parallelModel:add(tmodel)
parallelModel:add(nn.Identity())
parallelModel:add(nn.Identity())
model = parallelModel
if true then
local function sharingKey(m)
local key = torch.type(m)
if m.__shareGradInputKey then
key = key .. ':' .. m.__shareGradInputKey
end
return key
end
-- Share gradInput for memory efficient backprop
local cache = {}
model:apply(function(m)
local moduleType = torch.type(m)
if torch.isTensor(m.gradInput) and moduleType ~= 'nn.ConcatTable' then
local key = sharingKey(m)
if cache[key] == nil then
cache[key] = torch.CudaStorage(1)
end
m.gradInput = torch.CudaTensor(cache[key], 1, 0)
end
end)
end
if true then
cudnn.fastest = true
cudnn.benchmark = true
-- Wrap the model with DataParallelTable, if using more than one GPU
local gpus = torch.range(1, 2):totable()
local fastest, benchmark = cudnn.fastest, cudnn.benchmark
local dpt = nn.DataParallelTable(1, true, true)
:add(model, gpus)
:threads(function()
local cudnn = require 'cudnn'
cudnn.fastest, cudnn.benchmark = fastest, benchmark
end)
dpt.gradInput = nil
model = dpt:cuda()
end
newmodel = nn.Sequential()
newmodel:add(model)
input1 = nn.Identity()()
input2 = nn.Identity()()
input3 = nn.Identity()()
out = newmodel({input1,input2,input3})
r1 = nn.NarrowTable(1,2)(out)
r2 = nn.NarrowTable(2,2)(out)
f1 = nn.JoinTable(2)(r1)
f2 = nn.JoinTable(2)(r2)
n1 = nn.Sequential()
n1:add(nn.Linear(20,5))
n2 = nn.Sequential()
n2:add(nn.Linear(20,5))
f11 = n1(f1)
f12 = n2(f2)
foutput = nn.JoinTable(2)({f11,f12})
g = nn.gModule({input1,input2,input3},{foutput})
g = g:cuda()
g:forward({data1, data2, data2})
g:backward({data1, data2, data2}, torch.rand(4,10):cuda())
The code in the "if" statements is taken from Facebook's ResNet implementation