FIX: RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu) and RuntimeError: result type Float can't be cast to the desired output type long int#284