Closed nix-apollo closed 8 months ago
The following fails with a device error:
from rib.loader import load_sequential_transformer import torch model, _ = load_sequential_transformer( node_layers=['mlp_in.5', 'mlp_out.5'], last_pos_module_type=None, tlens_pretrained="pythia-14m", tlens_model_path=None, eps=None, fold_bias=True, device="cuda", ) x = torch.tensor([[100, 200]], device="cuda") # model(x)
Fails due to a device error. Running model.to("cuda") fixes it.
model.to("cuda")
@nix-apollo I can't replicate this bug. Note that I had to remove the eps argument. Can you try again and close issue if it still exists.
Yep works fine for me now. Unsure what changed.
The following fails with a device error:
Fails due to a device error. Running
model.to("cuda")
fixes it.