Aspects described herein provide a method including: receiving input data at a machine learning model, comprising: a plurality of processing layers; a plurality of gate logics; a plurality of gates; and a fully connected layer; determining based on a plurality of gate parameters associated with the plurality of gate logics, a subset of the plurality of processing layers with which to process the input data; processing the input data with the subset of the plurality of processing layers and the fully connected layer to generate an inference; determining a prediction loss based on the inference and a training label associated with the input data; determining an energy loss based on the subset of the plurality of processing layers used to process the input data; and optimizing the machine learning model based on: the prediction loss; the energy loss; and a prior probability associated with the training label.