         while true {
             let observationPython = Tensor<Double>(numpy: observationNumpy).unwrapped()
             let actionProbabilities =
-              softmax(net.applied(to: Tensor(observationPython).reshaped(to: [1, 4]), in: context))
+              softmax(net.applied(to: Tensor(observationPython).reshaped(to: [1, 4])))
             let actionProbabilitiesPython = actionProbabilities[0].makeNumpyArray()
             let len = Python.len(actionProbabilitiesPython)
             assert(actionCount == Int(Python.len(actionProbabilitiesPython)))

) -> [Episode] {
    var observationNumpy = env.reset()

    let context = Context(learningPhase: .inference)
    var episodes: [Episode] = []

    // We build up a batch of observations and actions.

    }

    @differentiable
    func applied(to input: Tensor<Float>) -> Tensor<Float> {
        return input.sequenced(through: l1, l2)
    }
}

