 ) -> [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>, in context: Context) -> Tensor<Float> {
        return l2.applied(to: l1.applied(to: input, in: context), in: context)
    }
}


      episodes: episodes, actionCount: actionCount)

    let gradients = gradient(at: net) { model -> Tensor<Float> in
        let logits = model.applied(to: input)
        let loss = softmaxCrossEntropy(logits: logits, probabilities: target)
        print("loss is \(loss)")
        return loss
    }
