tcm03 commited on
Commit
e9e7244
·
unverified ·
1 Parent(s): 633bded

Add PyTorchModelHubMixin to CLIP

Browse files
Files changed (1) hide show
  1. code/clip/model.py +2 -1
code/clip/model.py CHANGED
@@ -7,6 +7,7 @@ import numpy as np
7
  import torch
8
  import torch.nn.functional as F
9
  from torch import nn
 
10
 
11
 
12
  class Bottleneck(nn.Module):
@@ -237,7 +238,7 @@ class VisualTransformer(nn.Module):
237
 
238
  return x
239
 
240
- class CLIP(nn.Module):
241
  def __init__(self,
242
  embed_dim: int,
243
  # vision
 
7
  import torch
8
  import torch.nn.functional as F
9
  from torch import nn
10
+ from huggingface_hub import PyTorchModelHubMixin
11
 
12
 
13
  class Bottleneck(nn.Module):
 
238
 
239
  return x
240
 
241
+ class CLIP(nn.Module, PyTorchModelHubMixin):
242
  def __init__(self,
243
  embed_dim: int,
244
  # vision