<reponame>z-a-f/zaf_funcs<filename>nn.py
import torch
import torch.nn as nn

class Flatten(nn.Module):
  """PyTorch layer to flatten the inputs."""
  def forward(self, x):
    return x.view(x.size()[0], -1)
import torch
import torch.nn as nn

class Flatten(nn.Module):
  """PyTorch layer to flatten the inputs."""
  def forward(self, x):
    return x.view(x.size()[0], -1)
PyTorch layer to flatten the inputs.