Update model/openllama.py

model/openllama.py

@@ -632,8 +632,10 @@ class OpenLLAMAPEFTModel(nn.Module):
                 anomaly_map = (100.0 * patch_tokens[layer] @ feats_text_tensor.transpose(-2,-1))
                 B, L, C = anomaly_map.shape
                 H = int(np.sqrt(L))
+                anomaly_map = anomaly_map.to(torch.float16)
                 anomaly_map = F.interpolate(anomaly_map.permute(0, 2, 1).view(B, 2, H, H),
                                             size=224, mode='bilinear', align_corners=True)
+                anomaly_map = anomaly_map.to(torch.bfloat16)
                 anomaly_map = torch.softmax(anomaly_map, dim=1)
                 anomaly_maps.append(anomaly_map[:,1,:,:])
 
@@ -657,7 +659,9 @@ class OpenLLAMAPEFTModel(nn.Module):
                     sims.append(sim_max)
 
                 sim = torch.mean(torch.stack(sims,dim=0), dim=0).reshape(1,1,16,16)
+                anomaly_map = anomaly_map.to(torch.float16)
                 sim = F.interpolate(sim,size=224, mode='bilinear', align_corners=True)
+                anomaly_map = anomaly_map.to(torch.bfloat16)
                 anomaly_map_ret = 1 - sim # (anomaly_map_ret + 1 - sim) / 2