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 {
2222 public var dense3 = Dense<Float>(inputSize: 4, outputSize: 4)
2323 public var flatten = Flatten<Float>()
2424
25- @differentiable
26- public func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
25+ @differentiable(reverse) public func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
2726 let layer1 = dense1(input)
2827 let layer2 = layer1.reshaped(to: [1, 4])
2928 let layer3 = dense2(layer2)
Original file line number Diff line number Diff line change @@ -34,8 +34,7 @@ struct Model: Layer {
3434 var layer2 = Dense< Float > (inputSize : hiddenSize, outputSize : hiddenSize, activation : relu)
3535 var layer3 = Dense< Float > (inputSize : hiddenSize, outputSize : 3 , activation : identity)
3636
37- @differentiable
38- func callAsFunction (_ input : Tensor<Float >) -> Tensor<Float > {
37+ @differentiable (reverse) func callAsFunction (_ input : Tensor<Float >) -> Tensor<Float > {
3938 return input.sequenced (through : layer1, layer2, layer3)
4039 }
4140}
You can’t perform that action at this time.
0 commit comments