:mod:`unet`
===========

.. py:module:: unet

.. autoapi-nested-parse::

   Original U-Net model implementation.



Module Contents
---------------

.. py:class:: Model(hparams)

   Bases: :class:`model.base_model.BaseModel`

   .. autoapi-inheritance-diagram:: unet.Model
      :parts: 1

   This class implements U-Net module and its training functionality. It is equivalent
       to PyTorch's nn.Module in all aspects.

   :param LightningModule: The Pytorch-Lightning module derived from nn.module with
       useful hooks
   :type LightningModule: nn.Module

   .. method:: forward(self, x)


      Does the forward pass on the model.

      :param x: Input tensor batch.
      :type x: torch.Tensor
      :return: Output activations.
      :rtype: torch.Tensor