o

    /%h                     @   s<   d dl Z d dlmZ d dlm  mZ G dd dejZdS )    Nc                       s&   e Zd Zd fdd	Zdd Z  ZS )ContrastiveHead      c              	      sF   t    tt||t|tjddt||t|| _d S )NT)inplace)super__init__nn
SequentialLinearBatchNorm1dReLUmlp)selfin_dimZproj_dim	__class__ 9/root/Awesome-Backbones/configs/necks/contrastive_head.pyr      s   




zContrastiveHead.__init__c                 C   s   t j| |ddS )N   )dim)F	normalizer
   )r   xr   r   r   forward   s   zContrastiveHead.forward)r   r   )__name__
__module____qualname__r   r   
__classcell__r   r   r   r   r      s    	r   )torchtorch.nnr   Ztorch.nn.functional
functionalr   Moduler   r   r   r   r   <module>   s    