From 711b01bd061a892aa44294fbad703b2ceddd0478 Mon Sep 17 00:00:00 2001 From: Mohamed Ezz Date: Tue, 22 Aug 2017 15:11:36 +0200 Subject: [PATCH] Handle soundnet5 case where batchnorm layer is cudnn.SpatialBatchNormalization --- load_t7.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/load_t7.py b/load_t7.py index 3665d3c..954d9fb 100644 --- a/load_t7.py +++ b/load_t7.py @@ -41,7 +41,8 @@ def load(o, param_list): param_list.append(temp) # batch norm elif o['modules'][i]._typename == 'nn.SpatialBatchNormalization' or \ - o['modules'][i]._typename == 'nn.VolumetricBatchNormalization': + o['modules'][i]._typename == 'cudnn.SpatialBatchNormalization' or \ + o['modules'][i]._typename == 'nn.VolumetricBatchNormalization': param_list[-1]['gamma'] = o['modules'][i]['weight'] param_list[-1]['beta'] = o['modules'][i]['bias'] param_list[-1]['mean'] = o['modules'][i]['running_mean']