fix: batch norm weights tensor to Dense

i8run · Feb 9, 2018 · 8349eb4 · 8349eb4
1 parent 97a9b41
commit 8349eb4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/mkldnn/Linear.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/mkldnn/Linear.scala
@@ -151,7 +151,7 @@ class Linear[T: ClassTag](
     primitive
   }
 
-  var _shouldConvert: Boolean = true
+  var _shouldConvert: Boolean = false
   def shouldConvert: Boolean = _shouldConvert
   def setShouldConvert(v: Boolean): this.type = {
     _shouldConvert = v
@@ -683,7 +683,7 @@ class Linear[T: ClassTag](
 
     if (gradOutputReorderPrim != 0) {
       if (gradOutput.getTensorType == DenseType && gradOutputPtr != 0) {
-        MklDnn.MemoryReleaseDataHandle(weight.storage().array().asInstanceOf[Array[Float]],
+        MklDnn.MemoryReleaseDataHandle(gradOutput.storage().array().asInstanceOf[Array[Float]],
           gradOutputPtr)
       }
     }

diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/mkldnn/SpatialBatchNormalization.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/mkldnn/SpatialBatchNormalization.scala
@@ -417,7 +417,7 @@ class SpatialBatchNormalization[T: ClassTag](
   // in mkl dnn, the weight and bias should be all in the same array
   private def createParams(initWeight: Tensor[T], initBias: Tensor[T]): Tensor[T] = {
     val weightAndBias: Tensor[T] = if (affine) {
-      MklDnnTensor[T](Array(2 * nOutput))
+      Tensor[T](Array(2 * nOutput))
     } else {
       null
     }