credit.postblock.wet_mask_samudra
=================================

.. py:module:: credit.postblock.wet_mask_samudra


Classes
-------

.. autoapisummary::

   credit.postblock.wet_mask_samudra.WetMaskBlock


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

.. py:class:: WetMaskBlock(conf)

   Bases: :py:obj:`torch.nn.Module`


   Post-processing layer that applies wet mask to ocean predictions
   Zero trainable parameters, but mask influences gradients

   Masks predictions so land points = 0, ocean points preserve values.
   This encourages the model to focus learning on ocean regions.


   .. py:method:: forward(predictions)

      Apply wet mask to predictions with gradient influence

      :param predictions: tensor of shape (batch, n_vars, time, lat, lon)

      :returns: same shape, with land values set to zero
      :rtype: masked_predictions