Batch operation problem for specific conditions

Hi guys

I’m doing a conditional sampling test from the results of the network

For example,

cond = torch.where(x> 0)

If the condition condition is not satisfied in a specific batch, 0 data is sampled from x and the mean is calculated, resulting in nan.

The solution, in my opinion, is to give each condition to the batch while running the for loop. Can I prevent this problem with other simple operations?

I think checking for an empty output before applying the mean operation sounds like a simple and valid solution.
What kind of solution are you looking for and would the check not work for you?