I was able to update the gist and make it even shorter, with no dependence on the higher library

  
      gist.github.com
  
  
    https://gist.github.com/jeffwillette/3d48fd530c456311613c390a2aa0dfae 10
gistfile1.txt
from argparse import Namespace
from typing import Any

import torch
from torch import nn
from torch.nn import functional as F

T = torch.Tensor


This file has been truncated. show original