Fix FusedLayerNorm with CPU inputs
Summary: Pull Request resolved:
@@ -16,6 +16,9 @@ try:
class FusedLayerNorm(_FusedLayerNorm):
@torch.jit.unused
def forward(self, x):
+ if not x.is_cuda:
+ return super().forward(x)
+ else:
with torch.cuda.device(x.device):
return super().forward(x)
