examples/research_projects/deebert/src/modeling_highway_bert.py/HighwayException/__init__
class HighwayException:    def __init__(self, message, exit_layer):
        self.message = message
        self.exit_layer = exit_layer  # start from 1!