use enable in _vjpCausallyMasked (#218)

koen-dejonghe · rxwei · commit ae8eac0677ed · 2019-10-19T11:31:34.000-07:00
Fixes #217
diff --git a/Transformer/Model.swift b/Transformer/Model.swift
@@ -102,7 +102,7 @@ func causallyMasked(_ dotProducts: Tensor<Float>, enable: Bool = false) -> Tenso
 // causal mask is intentionally invisible to differentiation
 func _vjpCausallyMasked(_ dotProducts: Tensor<Float>, enable: Bool)
     -> (Tensor<Float>, (Tensor<Float>) -> Tensor<Float>) {
-    return (causallyMasked(dotProducts), identity)
+    return (causallyMasked(dotProducts, enable: enable), identity)
 }
 
 struct Attention: ParameterlessLayer {

Original file line number	Diff line number	Diff line change
`@@ -102,7 +102,7 @@ func causallyMasked(_ dotProducts: Tensor<Float>, enable: Bool = false) -> Tenso`
`102`	`102`	`// causal mask is intentionally invisible to differentiation`
`103`	`103`	`func _vjpCausallyMasked(_ dotProducts: Tensor<Float>, enable: Bool)`
`104`	`104`	`-> (Tensor<Float>, (Tensor<Float>) -> Tensor<Float>) {`
`105`		`- return (causallyMasked(dotProducts), identity)`
	`105`	`+ return (causallyMasked(dotProducts, enable: enable), identity)`
`106`	`106`	`}`
`107`	`107`
`108`	`108`	`struct Attention: ParameterlessLayer {`