One missing value

x = torch.tensor([1., float("nan"), 3.])

x.sum()       # nan
x.mean()      # nan
x.nansum()    # 4.0
x.nanmean()   # 2.0

nan propagates through arithmetic by definition, so one of them anywhere in a tensor makes every reduction over it nan. That is the IEEE behaviour and it is what you want by default: a silent nan is how a training run dies at step 40,000 with no explanation, and propagation is what lets you find it.

When to reach for the nan-aware version

When nan genuinely means “missing”, not “something went wrong”. Sensor data with dropouts, ragged sequences padded with nan, metrics where some examples had no valid target.

x.nansum()    # treats nan as 0
x.nanmean()   # ignores nan in both the sum AND the count

That second point is the one worth noticing. nanmean of [1, nan, 3] is 2.0, not 1.333: it divides by 2, the number of real values, not by 3. Substituting zeros yourself and taking a plain mean gets that wrong.

When not to

If nan means your loss exploded, silencing it with nansum converts a loud failure into a model that trains on garbage. Reach for torch.isnan(x).any() and an assertion instead.

The distinction is whether the nan is data or a symptom. Only you know that, which is why PyTorch does not choose.

Your task

def mean_of_valid(x: torch.Tensor) -> torch.Tensor

Return the mean of the non-nan entries of x, over the last axis, as a tensor. A row that is entirely nan should give nan.