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]

Codecov Report

All modified and coverable lines are covered by tests ✅

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

Additional details and impacted files

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

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[weihua916]

Thanks!