Don't put assertions on N/A imputation correctness (#433)

I understand the intention of the assertion introduced in #307, however,
we should not check it at runtime as it causes a D2H sync. It should
rather be covered in tests, and I believe there's no need to add new
tests as existing test cases already cover the correctness of the N/A
imputation.

CHANGELOG.md

@@ -18,7 +18,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
-- Removed CUDA synchronizations in `nn.LinearEmbeddingEncoder` ([#432](https://github.com/pyg-team/pytorch-frame/pull/432))
+- Removed CUDA synchronizations from `nn.LinearEmbeddingEncoder` ([#432](https://github.com/pyg-team/pytorch-frame/pull/432))
+- Removed CUDA synchronizations from N/A imputation logic in `nn.StypeEncoder` ([#433](https://github.com/pyg-team/pytorch-frame/pull/433))
 
 ## [0.2.3] - 2024-07-08
 

torch_frame/nn/encoder/stype_encoder.py

@@ -250,16 +250,7 @@ def na_forward(self, feat: TensorData) -> TensorData:
                 fill_values = torch.tensor(fill_values, device=feat.device)
                 assert feat.size(-1) == fill_values.size(-1)
                 feat = torch.where(na_mask, fill_values, feat)
-        # Add better safeguard here to make sure nans are actually
-        # replaced, expecially when nans are represented as -1's. They are
-        # very hard to catch as they won't error out.
-        filled_values = feat
-        if isinstance(feat, _MultiTensor):
-            filled_values = feat.values
-        if filled_values.is_floating_point():
-            assert not torch.isnan(filled_values).any()
-        else:
-            assert not (filled_values == -1).any()
+
         return feat