wengong-jin / dsmbind Goto Github PK

RuntimeError Traceback (most recent call last)
Cell In[12], line 1
----> 1 sabdab_evaluate(model, test_data, embedding, model_args)

Cell In[11], line 6, in sabdab_evaluate(model, data, embedding, args)
4 for ab in tqdm(data):
5 binder, target = AntibodyDataset.make_local_batch([ab], embedding, args)
----> 6 score = model.predict(binder, target)
7 pred.append(-1.0 * score.item())
8 label.append(ab['affinity'])

File ./bindenergy/models/energy.py:269, in AllAtomEnergyModel.predict(self, binder, target, visual)
266 D = (bind_X.view(B,N14,1,3) - tgt_X.view(B,1,M14,3)).norm(dim=-1) # [B,N14,M14]
267 mask_2D = mask_2D * (D < self.threshold).float()
--> 269 h = self.encoder(
270 (bind_X, bind_S, bind_A, None),
271 (tgt_X, tgt_S, tgt_A, None),
272 ) # [B,N+M,14,H]
274 bind_h = self.W_o(h[:, :N]).view(B, N14, -1)
275 tgt_h = self.U_o(h[:, N:]).view(B, M14, -1)

File ~/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
1516 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1517 else:
-> 1518 return self._call_impl(*args, **kwargs)

File ~/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
1522 # If we don't have any hooks, we want to skip the rest of the logic in
1523 # this function, and just call forward.
1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1525 or _global_backward_pre_hooks or _global_backward_hooks
1526 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527 return forward_call(*args, **kwargs)
1529 try:
1530 result = None

File ./bindenergy/models/frame.py:132, in AllAtomEncoder.forward(self, binder, target)
129 mask = mask.unsqueeze(1).expand(-1,8,-1).reshape(B*8, (N+M)*14)
131 h = torch.cat([h_X, h_S], dim=-1)
--> 132 h, _, _ = self.encoder(
133 h.transpose(0, 1),
134 mask_pad=(~mask.transpose(0, 1).bool())
135 )
136 h = h.transpose(0, 1).view(B, 8, (N+M)*14, -1)
137 h = h.mean(dim=1) # frame averaging

File ~/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
1516 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1517 else:
-> 1518 return self._call_impl(*args, **kwargs)

File ~/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
1522 # If we don't have any hooks, we want to skip the rest of the logic in
1523 # this function, and just call forward.
1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1525 or _global_backward_pre_hooks or _global_backward_hooks
1526 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527 return forward_call(*args, **kwargs)
1529 try:
1530 result = None

File ~/conda/lib/python3.8/site-packages/sru/modules.py:1215, in SRUpp.forward(self, input, c0, mask_pad, attn_mask, memory, memory_mask_pad)
1213 prev_inputs.append(x)
1214 memory_i = memory[i] if memory is not None else None
-> 1215 h, c = rnn(x, c0_[i],
1216 mask_pad=mask_pad,
1217 attn_mask=attn_mask,
1218 memory=memory_i,
1219 memory_mask_pad=memory_mask_pad)
1220 x = h
1221 lstc.append(c)

File ~/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
1516 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1517 else:
-> 1518 return self._call_impl(*args, **kwargs)

File ~/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
1522 # If we don't have any hooks, we want to skip the rest of the logic in
1523 # this function, and just call forward.
1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1525 or _global_backward_pre_hooks or _global_backward_hooks
1526 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527 return forward_call(*args, **kwargs)
1529 try:
1530 result = None

File ~/conda/lib/python3.8/site-packages/sru/modules.py:930, in SRUppCell.forward(self, input, c0, mask_pad, attn_mask, memory, memory_mask_pad)
927 V = self.weight_c
929 # apply elementwise recurrence to get hidden states h and c
--> 930 h, c = self.apply_recurrence(U, V,
931 residual, c0,
932 scale_val,
933 mask_c,
934 mask_pad)
936 layer_norm = self.layer_norm
937 if layer_norm is not None:

File ~/conda/lib/python3.8/site-packages/sru/modules.py:268, in SRUCell.apply_recurrence(self, U, V, residual, c0, scale_val, mask_c, mask_pad)
266 if not torch.jit.is_scripting():
267 if self.bias.is_cuda:
--> 268 return elementwise_recurrence_gpu(U, residual, V, self.bias, c0,
269 self.activation_type,
270 self.hidden_size,
271 self.bidirectional,
272 self.has_skip_term,
273 scale_val, mask_c, mask_pad,
274 self.amp_recurrence_fp16)
275 else:
276 return elementwise_recurrence_naive(U, residual, V, self.bias, c0,
277 self.activation_type,
278 self.hidden_size,
279 self.bidirectional,
280 self.has_skip_term,
281 scale_val, mask_c, mask_pad)

File ~/conda/lib/python3.8/site-packages/sru/ops.py:130, in elementwise_recurrence_gpu(U, x, weight_c, bias, c_init, activation_type, hidden_size, bidirectional, has_skip_term, scale_x, dropout_mask_c, mask_pad, amp_recurrence_fp16)
127 scale_x = cast(scale_x) if scale_x is not None else scale_x
128 dropout_mask_c = cast(dropout_mask_c) if dropout_mask_c is not None else dropout_mask_c
--> 130 return ElementwiseRecurrence.apply(
131 U,
132 x,
133 weight_c,
134 bias,
135 c_init,
136 activation_type,
137 hidden_size,
138 bidirectional,
139 has_skip_term,
140 scale_x,
141 dropout_mask_c,
142 mask_pad
143 )

File ~/conda/lib/python3.8/site-packages/torch/autograd/function.py:539, in Function.apply(cls, *args, **kwargs)
536 if not torch._C._are_functorch_transforms_active():
537 # See NOTE: [functorch vjp and autograd interaction]
538 args = _functorch.utils.unwrap_dead_wrappers(args)
--> 539 return super().apply(*args, **kwargs) # type: ignore[misc]
541 if cls.setup_context == _SingleLevelFunction.setup_context:
542 raise RuntimeError(
543 "In order to use an autograd.Function with functorch transforms "
544 "(vmap, grad, jvp, jacrev, ...), it must override the setup_context "
545 "staticmethod. For more details, please see "
546 "https://pytorch.org/docs/master/notes/extending.func.html"
547 )

File ~/conda/lib/python3.8/site-packages/sru/cuda_functional.py:176, in ElementwiseRecurrence.forward(ctx, u, x, weight_c, bias, init, activation_type, d_out, bidirectional, has_skip_term, scale_x, mask_c, mask_pad)
174 # ensure mask_pad is a bool tensor
175 mask_pad = mask_pad.bool().contiguous() if mask_pad is not None else None
--> 176 h, last_hidden, c = elementwise_recurrence_forward(
177 u,
178 x,
179 weight_c,
180 bias,
181 init,
182 activation_type,
183 d_out,
184 bidirectional,
185 has_skip_term,
186 scale_x,
187 mask_c,
188 mask_pad
189 )
190 ctx.save_for_backward(u, x, weight_c, bias, init, mask_c, c, mask_pad, scale_x)
191 return h, last_hidden

RuntimeError: Caught an unknown exception!