credit.postblock.wet_mask_samudra

Classes

WetMaskBlock

Post-processing layer that applies wet mask to ocean predictions

Module Contents

class credit.postblock.wet_mask_samudra.WetMaskBlock(conf)

Bases: 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.

forward(predictions)

Apply wet mask to predictions with gradient influence

Parameters:

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

Returns:

same shape, with land values set to zero

Return type:

masked_predictions