Extending the solution to support devices 
from typing import Optional

def roll(x: torch.Tensor, shift: int, dim: int = -1, fill_pad: Optional[int] = None):

    device = x.device
    
    if 0 == shift:
        return x

    elif shift < 0:
        shift = -shift
        gap = x.index_select(dim, torch.arange(shift, device=device))
        if fill_pad is not None:
            gap = fill_pad * torch.ones_like(gap, device=device)
        return torch.cat([x.index_select(dim, torch.arange(shift, x.size(dim), device=device)), gap], dim=dim)

    else:
        shift = x.size(dim) - shift
        gap = x.index_select(dim, torch.arange(shift, x.size(dim), device=device))
        if fill_pad is not None:
            gap = fill_pad * torch.ones_like(gap, device=device)
        return torch.cat([gap, x.index_select(dim, torch.arange(shift, device=device))], dim=dim)