The following snippet solves your problem:
import torch

a = torch.tensor(
[
    [0.5, 0.8, 0.9, 0.3],
    [0.5, 0.8, 0.7, 0.3],
    [0.6, 0.4, 0.3, 0.4],
])
r = 0.7

torch.min(a.max(1)[1], ((a > r).cumsum(1) < 1).sum(1))

Note that (at least for torch 1.1) torch.zeros(n).argmax() == n - 1 and torch.zeros(n).cuda().argmax() == 0… which are logically both valid, but aren’t consistent. If your operations only use cpu tensors, you can do it a little simpler with the following:
torch.min(a.max(1)[1], (a > r).cumsum(1).argmin(1) + 1)