fix(fcpe) : windows size 8 to 9
fix(fcpe) : windows size 8 to 9
This commit is contained in:
parent
d8c91d675b
commit
bb610c5805
|
@ -136,7 +136,7 @@ class FCPE(nn.Module):
|
||||||
B, N, _ = y.size()
|
B, N, _ = y.size()
|
||||||
ci = self.cent_table[None, None, :].expand(B, N, -1)
|
ci = self.cent_table[None, None, :].expand(B, N, -1)
|
||||||
confident, max_index = torch.max(y, dim=-1, keepdim=True)
|
confident, max_index = torch.max(y, dim=-1, keepdim=True)
|
||||||
local_argmax_index = torch.arange(0,8).to(max_index.device) + (max_index - 4)
|
local_argmax_index = torch.arange(0,9).to(max_index.device) + (max_index - 4)
|
||||||
local_argmax_index[local_argmax_index<0] = 0
|
local_argmax_index[local_argmax_index<0] = 0
|
||||||
local_argmax_index[local_argmax_index>=self.n_out] = self.n_out - 1
|
local_argmax_index[local_argmax_index>=self.n_out] = self.n_out - 1
|
||||||
ci_l = torch.gather(ci,-1,local_argmax_index)
|
ci_l = torch.gather(ci,-1,local_argmax_index)
|
||||||
|
|
Loading…
Reference in New Issue