@@ -437,6 +437,9 @@ class SwinBlock(nn.Module):
|
|||||||
hidden_states_windows = window_partition(shifted_hidden_states, self.window_size)
|
hidden_states_windows = window_partition(shifted_hidden_states, self.window_size)
|
||||||
hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels)
|
hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels)
|
||||||
|
|
||||||
|
if self.attn_mask is not None:
|
||||||
|
self.attn_mask = self.attn_mask.to(hidden_states_windows.device)
|
||||||
|
|
||||||
self_attention_outputs = self.attention(
|
self_attention_outputs = self.attention(
|
||||||
hidden_states_windows,
|
hidden_states_windows,
|
||||||
self.attn_mask,
|
self.attn_mask,
|
||||||
|
|||||||
@@ -595,8 +595,8 @@ class ViltModelIntegrationTest(unittest.TestCase):
|
|||||||
|
|
||||||
# forward pass
|
# forward pass
|
||||||
outputs = model(
|
outputs = model(
|
||||||
input_ids=encoding_1.input_ids,
|
input_ids=encoding_1.input_ids.to(torch_device),
|
||||||
pixel_values=pixel_values,
|
pixel_values=pixel_values.to(torch_device),
|
||||||
)
|
)
|
||||||
|
|
||||||
# verify the logits
|
# verify the logits
|
||||||
|
|||||||
@@ -327,9 +327,6 @@ class ViTMAEModelTest(ModelTesterMixin, unittest.TestCase):
|
|||||||
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
|
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
|
||||||
|
|
||||||
for model_class in self.all_model_classes:
|
for model_class in self.all_model_classes:
|
||||||
|
|
||||||
print("Model class:", model_class)
|
|
||||||
|
|
||||||
model = model_class(config)
|
model = model_class(config)
|
||||||
model.to(torch_device)
|
model.to(torch_device)
|
||||||
model.eval()
|
model.eval()
|
||||||
|
|||||||
Reference in New Issue
Block a user