Improve performance of Encoder's `na_forward` method

[DamianSzwichtenberg]

This PR introduces small improvements that resulted in an average 1.5x speedup in na_forward execution time. get_na_mask is a heavy operation, whose result may be cached at the pre-condition checking step. Additionally, when we deal with 2D feat, we can use torch.where to efficiently fill NaNs across all columns.

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 93.59%. Comparing base (2d166c6) to head (769e0f9).

Additional details and impacted files

```diff @@ Coverage Diff @@ ## master #364 +/- ## ======================================= Coverage 93.59% 93.59% ======================================= Files 119 119 Lines 6275 6280 +5 ======================================= + Hits 5873 5878 +5 Misses 402 402 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.