text: import torch
  from subspace_inference import unflatten_like
  
  class SubspaceModel(torch.nn.Module):
      def __init__(self, mean, cov_factor):
          super(SubspaceModel, self).__init__()
          self.rank = cov_factor.size(0)
          self.register_buffer('mean', mean)
          self.register_buffer('cov_factor', cov_factor)
  
      def forward(self, t):
          return self.mean + self.cov_factor.t() @ t
  
  class ProjectedModel(torch.nn.Module):
      def __init__(self, proj_params, model, projection=None, mean=None, subspace=None):
          super(ProjectedModel, self).__init__()
          self.model = model
  
          if subspace is None:
              self.subspace = SubspaceModel(mean, projection)
          else:
              self.subspace = subspace
  
          if mean is None and subspace is None:
              raise NotImplementedError('Must enter either subspace or mean')
  
          self.proj_params = proj_params
  
      def update_params(self, vec, model):
          vec_list = unflatten_like(likeTensorList=list(model.parameters()), vector=vec.view(1,-1))
          for param, v in zip(model.parameters(), vec_list):
              param.detach_()
              param.mul_(0.0).add_(v)
  
      def forward(self, *args, **kwargs):
          y = self.subspace(self.proj_params)
  
          self.update_params(y, self.model)
          return self.model(*args, **kwargs)
  