File tree Expand file tree Collapse file tree 2 files changed +2
-4
lines changed Expand file tree Collapse file tree 2 files changed +2
-4
lines changed Original file line number Diff line number Diff line change @@ -22,8 +22,7 @@ public struct MyModel: Layer {
22
22
public var dense3 = Dense<Float>(inputSize: 4, outputSize: 4)
23
23
public var flatten = Flatten<Float>()
24
24
25
- @differentiable
26
- public func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
25
+ @differentiable(reverse) public func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
27
26
let layer1 = dense1(input)
28
27
let layer2 = layer1.reshaped(to: [1, 4])
29
28
let layer3 = dense2(layer2)
Original file line number Diff line number Diff line change @@ -34,8 +34,7 @@ struct Model: Layer {
34
34
var layer2 = Dense< Float > (inputSize : hiddenSize, outputSize : hiddenSize, activation : relu)
35
35
var layer3 = Dense< Float > (inputSize : hiddenSize, outputSize : 3 , activation : identity)
36
36
37
- @differentiable
38
- func callAsFunction (_ input : Tensor<Float >) -> Tensor<Float > {
37
+ @differentiable (reverse) func callAsFunction (_ input : Tensor<Float >) -> Tensor<Float > {
39
38
return input.sequenced (through : layer1, layer2, layer3)
40
39
}
41
40
}
You can’t perform that action at this time.
0 commit comments