class_docstr: A deep energy module that unrolls an optimizer over the energy function.
  
  This module takes a grayscale 28x28 image x as the input and
  outputs a class prediction by (approximately) solving the
  optimization problem
  
      \hat y = argmin_y E_    heta(x, y)
  
  with a fixed number of gradient steps.
  
  Args:
      Enet: The energy network.
      n_cls (int): The number of classes.
      n_inner_iter (int): The number of optimization steps to take.
class_docstr_tok: ["A","deep","energy","module","that","unrolls","an","optimizer","over","the","energy","function",".","This","module","takes","a","grayscale","28x28","image","x","as","the","input","and","outputs","a","class","prediction","by","(","approximately",")","solving","the","optimization","problem","\\hat","y","=","argmin_y","E_","heta","(","x",",","y",")","with","a","fixed","number","of","gradient","steps",".","Args",":","Enet",":","The","energy","network",".","n_cls","(","int",")",":","The","number","of","classes",".","n_inner_iter","(","int",")",":","The","number","of","optimization","steps","to","take","."]
code: def __init__(
      self, Enet: EnergyNet, n_cls: int = 10, n_inner_iter: int = 5
  ):
      super(UnrollEnergy, self).__init__()
      self.Enet = Enet
      self.n_cls = n_cls
      self.n_inner_iter = n_inner_iter
  
imports: import argparse
  import typing
  import torch
  from torch import nn
  from torch.utils.data import DataLoader
  import torch.nn.functional
  import torch.optim
  from torchvision import datasets, transforms
  import higher
label_desc: Use image node_id 1 to create a new UnrollEnergy object from inherited base classes: nn with example: obj = UnrollEnergy(Enet, n_cls, n_inner_iter)